KT-13241 Indices optimization leads to VerifyError with smart cast receiver

Use expected receiver type from the corresponding extension function.

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -635,38 +635,31 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     ) {
         CallableDescriptor loopRangeCallee = loopRangeCall.getResultingDescriptor();
         if (RangeCodegenUtil.isOptimizableRangeTo(loopRangeCallee)) {
-            ReceiverValue from = loopRangeCall.getDispatchReceiver();
-            assert from != null : "Dispatch receiver should be non-null for optimizable 'rangeTo' call";
-            KtExpression to = getSingleArgumentExpression(loopRangeCall);
-            assert to != null : "Optimizable 'rangeTo' call should have a single expression argument";
-            return new ForInRangeLiteralLoopGenerator(forExpression, from, to);
+            return new ForInRangeLiteralLoopGenerator(forExpression, loopRangeCall);
         }
         else if (RangeCodegenUtil.isOptimizableDownTo(loopRangeCallee)) {
-            ReceiverValue from = loopRangeCall.getExtensionReceiver();
-            assert from != null : "Extension receiver should be non-null for optimizable 'downTo' call";
-            KtExpression to = getSingleArgumentExpression(loopRangeCall);
-            assert to != null : "Optimizable 'downTo' call should have a single expression argument";
-            return new ForInDownToProgressionLoopGenerator(forExpression, from, to);
+            return new ForInDownToProgressionLoopGenerator(forExpression, loopRangeCall);
         }
         else if (RangeCodegenUtil.isArrayOrPrimitiveArrayIndices(loopRangeCallee)) {
-            ReceiverValue extensionReceiver = loopRangeCall.getExtensionReceiver();
-            assert extensionReceiver != null : "Extension receiver should be non-null for optimizable 'Array.indices' call";
-            return new ForInArrayIndicesRangeLoopGenerator(forExpression, extensionReceiver);
+            return new ForInArrayIndicesRangeLoopGenerator(forExpression, loopRangeCall);
         }
         else if (RangeCodegenUtil.isCollectionIndices(loopRangeCallee)) {
-            ReceiverValue extensionReceiver = loopRangeCall.getExtensionReceiver();
-            assert extensionReceiver != null : "Extension receiver should be non-null for optimizable 'Collection.indices' call";
-            return new ForInCollectionIndicesRangeLoopGenerator(forExpression, extensionReceiver);
+            return new ForInCollectionIndicesRangeLoopGenerator(forExpression, loopRangeCall);
         }
         else if (RangeCodegenUtil.isCharSequenceIndices(loopRangeCallee)) {
-            ReceiverValue extensionReceiver = loopRangeCall.getExtensionReceiver();
-            assert extensionReceiver != null : "Extension receiver should be non-null for optimizable 'CharSequence.indices' call";
-            return new ForInCharSequenceIndicesRangeLoopGenerator(forExpression, extensionReceiver);
+            return new ForInCharSequenceIndicesRangeLoopGenerator(forExpression, loopRangeCall);
         }
 
         return null;
     }
 
+    @NotNull
+    private static KotlinType getExpectedReceiverType(@NotNull ResolvedCall<? extends CallableDescriptor> resolvedCall) {
+        ReceiverParameterDescriptor extensionReceiver = resolvedCall.getResultingDescriptor().getExtensionReceiverParameter();
+        assert extensionReceiver != null : "Extension receiver should be non-null";
+        return extensionReceiver.getType();
+    }
+
     @Nullable
     private static KtExpression getSingleArgumentExpression(@NotNull ResolvedCall<? extends CallableDescriptor> resolvedCall) {
         List<ResolvedValueArgument> resolvedValueArguments = resolvedCall.getValueArgumentsByIndex();
@@ -1134,14 +1127,10 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         private final ReceiverValue from;
         private final KtExpression to;
 
-        private ForInRangeLiteralLoopGenerator(
-                @NotNull KtForExpression forExpression,
-                @NotNull ReceiverValue from,
-                @NotNull KtExpression to
-        ) {
+        private ForInRangeLiteralLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ResolvedCall<?> loopRangeCall) {
             super(forExpression);
-            this.from = from;
-            this.to = to;
+            this.from = loopRangeCall.getDispatchReceiver();
+            this.to = getSingleArgumentExpression(loopRangeCall);
         }
 
         @Override
@@ -1155,14 +1144,10 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         private final ReceiverValue from;
         private final KtExpression to;
 
-        private ForInDownToProgressionLoopGenerator(
-                @NotNull KtForExpression forExpression,
-                @NotNull ReceiverValue from,
-                @NotNull KtExpression to
-        ) {
+        private ForInDownToProgressionLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ResolvedCall<?> loopRangeCall) {
             super(forExpression, -1);
-            this.from = from;
-            this.to = to;
+            this.from = loopRangeCall.getExtensionReceiver();
+            this.to = getSingleArgumentExpression(loopRangeCall);
         }
 
         @Override
@@ -1193,10 +1178,12 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
     private abstract class ForInOptimizedIndicesLoopGenerator extends AbstractForInRangeLoopGenerator {
         protected final ReceiverValue receiverValue;
+        protected final KotlinType expectedReceiverType;
 
-        private ForInOptimizedIndicesLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
+        private ForInOptimizedIndicesLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ResolvedCall<?> loopRangeCall) {
             super(forExpression);
-            this.receiverValue = receiverValue;
+            this.receiverValue = loopRangeCall.getExtensionReceiver();
+            this.expectedReceiverType = getExpectedReceiverType(loopRangeCall);
         }
 
         @Override
@@ -1204,7 +1191,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             loopParameter().store(StackValue.constant(0, asmElementType), v);
 
             StackValue receiver = generateReceiverValue(receiverValue, false);
-            Type receiverType = asmType(receiverValue.getType());
+            Type receiverType = asmType(expectedReceiverType);
             receiver.put(receiverType, v);
             getReceiverSizeAsInt();
             v.iconst(1);
@@ -1219,8 +1206,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     private class ForInCollectionIndicesRangeLoopGenerator extends ForInOptimizedIndicesLoopGenerator {
-        private ForInCollectionIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
-            super(forExpression, receiverValue);
+        private ForInCollectionIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ResolvedCall<?> loopRangeCall) {
+            super(forExpression, loopRangeCall);
         }
 
         @Override
@@ -1230,8 +1217,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     private class ForInArrayIndicesRangeLoopGenerator extends ForInOptimizedIndicesLoopGenerator {
-        private ForInArrayIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
-            super(forExpression, receiverValue);
+        private ForInArrayIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ResolvedCall<?> loopRangeCall) {
+            super(forExpression, loopRangeCall);
         }
 
         @Override
@@ -1241,8 +1228,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     private class ForInCharSequenceIndicesRangeLoopGenerator extends ForInOptimizedIndicesLoopGenerator {
-        private ForInCharSequenceIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
-            super(forExpression, receiverValue);
+        private ForInCharSequenceIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ResolvedCall<?> loopRangeCall) {
+            super(forExpression, loopRangeCall);
         }
 
         @Override

compiler/testData/codegen/box/ranges/forInIndices/kt13241_Array.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun test(x: Any): Int {
+    var sum = 0
+    if (x is IntArray) {
+        for (i in x.indices) {
+            sum = sum * 10 + i
+        }
+    }
+    return sum
+}
+
+fun box(): String {
+    assertEquals(123, test(intArrayOf(0, 0, 0, 0)))
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/kt13241_CharSequence.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun test(x: Any): Int {
+    var sum = 0
+    if (x is String) {
+        for (i in x.indices) {
+            sum = sum * 10 + i
+        }
+    }
+    return sum
+}
+
+fun box(): String {
+    assertEquals(123, test("0000"))
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/kt13241_Collection.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun test(x: Any): Int {
+    var sum = 0
+    if (x is List<*>) {
+        for (i in x.indices) {
+            sum = sum * 10 + i
+        }
+    }
+    return sum
+}
+
+fun box(): String {
+    assertEquals(123, test(listOf(0, 0, 0, 0)))
+    return "OK"
+}

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

@@ -11273,6 +11273,24 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/kt12983_forInSpecificCollectionIndices.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("kt13241_Array.kt")
+            public void testKt13241_Array() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/kt13241_Array.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt13241_CharSequence.kt")
+            public void testKt13241_CharSequence() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/kt13241_CharSequence.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt13241_Collection.kt")
+            public void testKt13241_Collection() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/kt13241_Collection.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/ranges/literal")