KT-5044 intrinsify all contains calls for primitive ranges from stdlib

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

@@ -3448,8 +3448,9 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         return StackValue.operation(Type.BOOLEAN_TYPE, new Function1<InstructionAdapter, Unit>() {
             @Override
             public Unit invoke(InstructionAdapter v) {
-                if (isIntRangeExpr(deparenthesized) && AsmUtil.isIntPrimitive(leftValue.type)) {
+                if (RangeCodegenUtil.isOptimizablePrimitiveRangeSpecialization(leftValue.type, deparenthesized, bindingContext)
-                    genInIntRange(leftValue, (KtBinaryExpression) deparenthesized, isInverted);
+                    || RangeCodegenUtil.isOptimizableRangeTo(operationReference, bindingContext)) {
+                    generateInPrimitiveRange(leftValue, (KtBinaryExpression) deparenthesized, isInverted);
                 }
                 else {
                     ResolvedCall<? extends CallableDescriptor> resolvedCall = CallUtilKt
@@ -3466,36 +3467,58 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     /*
-     * Translates x in a..b (for int and char ranges only) to a <= x && x <= b
+     * Translates x in a..b to a <= x && x <= b
-     * and x !in a..b to a > x || x > b
+     * and x !in a..b to a > x || x > b for any primitive type
      */
-    private void genInIntRange(StackValue leftValue, KtBinaryExpression rangeExpression, boolean isInverted) {
+    private void generateInPrimitiveRange(StackValue leftValue, KtBinaryExpression rangeExpression, boolean isInverted) {
-        int localVarIndex = myFrameMap.enterTemp(Type.INT_TYPE);
+        Type rangeType = leftValue.type;
-
+        int localVarIndex = myFrameMap.enterTemp(rangeType);
         // Load left bound
-        gen(rangeExpression.getLeft(), Type.INT_TYPE);
+        gen(rangeExpression.getLeft(), rangeType);
         // Load x into local variable to avoid StackValue#put side-effects
-        leftValue.put(Type.INT_TYPE, v);
+        leftValue.put(rangeType, v);
-        v.store(localVarIndex, Type.INT_TYPE);
+        v.store(localVarIndex, rangeType);
-        v.load(localVarIndex, Type.INT_TYPE);
+        v.load(localVarIndex, rangeType);
 
         // If (x < left) goto L1
         Label l1 = new Label();
-        v.ificmpgt(l1);
+        emitGreaterThan(rangeType, l1);
 
         // If (x > right) goto L1
-        v.load(localVarIndex, Type.INT_TYPE);
+        v.load(localVarIndex, rangeType);
-        gen(rangeExpression.getRight(), Type.INT_TYPE);
+        gen(rangeExpression.getRight(), rangeType);
-        v.ificmpgt(l1);
+        emitGreaterThan(rangeType, l1);
 
         Label l2 = new Label();
         v.iconst(isInverted ? 0 : 1);
         v.goTo(l2);
 
         v.mark(l1);
-        v.iconst(isInverted? 1 : 0);
+        v.iconst(isInverted ? 1 : 0);
         v.mark(l2);
-        myFrameMap.leaveTemp(Type.INT_TYPE);
+        myFrameMap.leaveTemp(rangeType);
+    }
+
+    private void emitGreaterThan(Type type, Label label) {
+        if (AsmUtil.isIntPrimitive(type)) {
+            v.ificmpgt(label);
+        }
+        else if (type == Type.LONG_TYPE) {
+            v.lcmp();
+            v.ifgt(label);
+        }
+        // '>' != 'compareTo' for NaN and +/- 0.0
+        else if (type == Type.FLOAT_TYPE) {
+            v.invokestatic("java/lang/Float", "compare", "(FF)I", false);
+            v.ifgt(label);
+        }
+        else if (type == Type.DOUBLE_TYPE) {
+            v.invokestatic("java/lang/Double", "compare", "(DD)I", false);
+            v.ifgt(label);
+        }
+        else {
+            throw new UnsupportedOperationException("Unexpected type: " + type);
+        }
     }
 
     private StackValue generateBooleanAnd(KtBinaryExpression expression) {
@@ -4540,19 +4563,6 @@ The "returned" value of try expression with no finally is either the last expres
         }
     }
 
-    private boolean isIntRangeExpr(KtExpression rangeExpression) {
-        if (rangeExpression instanceof KtBinaryExpression) {
-            KtBinaryExpression binaryExpression = (KtBinaryExpression) rangeExpression;
-            if (binaryExpression.getOperationReference().getReferencedNameElementType() == KtTokens.RANGE) {
-                KotlinType jetType = bindingContext.getType(rangeExpression);
-                assert jetType != null;
-                DeclarationDescriptor descriptor = jetType.getConstructor().getDeclarationDescriptor();
-                return DescriptorUtilsKt.getBuiltIns(descriptor).getIntegralRanges().contains(descriptor);
-            }
-        }
-        return false;
-    }
-
     private Call makeFakeCall(ReceiverValue initializerAsReceiver) {
         KtSimpleNameExpression fake = KtPsiFactoryKt.KtPsiFactory(state.getProject()).createSimpleName("fake");
         return CallMaker.makeCall(fake, initializerAsReceiver);

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

@@ -22,6 +22,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.builtins.PrimitiveType;
 import org.jetbrains.kotlin.descriptors.*;
+import org.jetbrains.kotlin.lexer.KtTokens;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.FqNameUnsafe;
 import org.jetbrains.kotlin.name.Name;
@@ -30,7 +31,11 @@ import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilKt;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
+import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver;
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
 import org.jetbrains.kotlin.types.KotlinType;
+import org.jetbrains.org.objectweb.asm.Type;
 
 import java.util.Arrays;
 import java.util.List;
@@ -123,6 +128,9 @@ public class RangeCodegenUtil {
         return getPrimitiveRangeOrProgressionElementType(className) != null;
     }
 
+    /*
+     * Checks whether rangeTo expression is optimizable for loop
+     */
     public static boolean isOptimizableRangeTo(CallableDescriptor rangeTo) {
         if ("rangeTo".equals(rangeTo.getName().asString())) {
             if (isPrimitiveNumberClassDescriptor(rangeTo.getContainingDeclaration())) {
@@ -176,6 +184,78 @@ public class RangeCodegenUtil {
         return true;
     }
 
+    /*
+     * Checks whether rangeTo expression is optimizable target of contains operator
+     */
+    public static boolean isOptimizableRangeTo(@NotNull KtSimpleNameExpression operationReference, @NotNull BindingContext bindingContext) {
+        ResolvedCall<? extends CallableDescriptor> resolvedCall = CallUtilKt
+                .getResolvedCallWithAssert(operationReference, bindingContext);
+        ReceiverValue receiver = resolvedCall.getDispatchReceiver();
+
+        /*
+         * Range is optimizable if
+         * 'in' receiver is expression 'rangeTo' from stdlib package and its argument
+         * has same primitive type as generic range parameter.
+         * For non-matching primitive types (e.g. int in double range)
+         * dispatch receiver will be null, because extension method will be called.
+         */
+        if (receiver instanceof ExpressionReceiver) {
+            ExpressionReceiver e = (ExpressionReceiver) receiver;
+            ResolvedCall<? extends CallableDescriptor> resolvedReceiver =
+                    CallUtilKt.getResolvedCall(e.getExpression(), bindingContext);
+
+            if (resolvedReceiver == null) {
+                return false;
+            }
+
+            CallableDescriptor descriptor = resolvedReceiver.getResultingDescriptor();
+            // kotlin.ranges.Ranges#rangeTo: ClosedRange<T> and T is primitive
+            // noinspection ConstantConditions
+            return isBuiltInRangeTo(descriptor) && KotlinBuiltIns.isPrimitiveType(descriptor.getExtensionReceiverParameter().getType());
+        }
+
+        return false;
+    }
+
+    private static boolean isBuiltInRangeTo(@NotNull CallableDescriptor descriptor) {
+        if (!isTopLevelInPackage(descriptor, "rangeTo", "kotlin.ranges")) {
+            return false;
+        }
+
+        ReceiverParameterDescriptor extensionReceiver = descriptor.getExtensionReceiverParameter();
+        return extensionReceiver != null;
+    }
+
+    /*
+     * Checks whether for expression 'x in a..b' a..b is primitive integral range
+     * with same type as x.
+     */
+    public static boolean isOptimizablePrimitiveRangeSpecialization(
+            @NotNull Type argumentType,
+            @NotNull KtExpression rangeExpression,
+            @NotNull BindingContext bindingContext
+    ) {
+        if (rangeExpression instanceof KtBinaryExpression) {
+            KtBinaryExpression binaryExpression = (KtBinaryExpression) rangeExpression;
+            if (binaryExpression.getOperationReference().getReferencedNameElementType() == KtTokens.RANGE) {
+                KotlinType kotlinType = bindingContext.getType(rangeExpression);
+                assert kotlinType != null;
+                DeclarationDescriptor descriptor = kotlinType.getConstructor().getDeclarationDescriptor();
+
+                // noinspection ConstantConditions
+                if (DescriptorUtilsKt.getBuiltIns(descriptor).getIntegralRanges().contains(descriptor)) {
+                    if ("LongRange".equals(descriptor.getName().asString())) {
+                        return argumentType == Type.LONG_TYPE;
+                    }
+
+                    return AsmUtil.isIntPrimitive(argumentType);
+                }
+            }
+        }
+
+        return false;
+    }
+
     private static boolean isTopLevelInPackage(@NotNull CallableDescriptor descriptor, @NotNull String name, @NotNull String packageName) {
         if (!name.equals(descriptor.getName().asString())) return false;
 

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

@@ -0,0 +1,25 @@
+// WITH_RUNTIME
+
+fun inInt(x: Long): Boolean {
+    return x in 1..2
+}
+
+fun inLong(x: Int): Boolean {
+    return x in 1L..2L
+}
+
+fun inFloat(x: Double): Boolean {
+    return x in 1.0f..2.0f
+}
+
+fun inDouble(x: Float): Boolean {
+    return x in 1.0..2.0
+}
+
+// 2 INVOKESPECIAL
+// 2 NEW
+// 2 INVOKESTATIC kotlin/ranges/RangesKt.rangeTo
+// 1 longRangeContains
+// 1 intRangeContains
+// 1 doubleRangeContains
+// 1 floatRangeContains

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

@@ -0,0 +1,33 @@
+// WITH_RUNTIME
+
+fun Byte.in2(left: Byte, right: Byte): Boolean {
+    return this in left..right
+}
+
+fun inInt(x: Int, left: Int, right: Int): Boolean {
+    return x in left..right
+}
+
+fun inDouble(x: Double, left: Double, right: Double): Boolean {
+    return x in left..right
+}
+
+fun inFloat(x: Float, left: Float, right: Float): Boolean {
+    return x in left..right
+}
+
+fun inLong(x: Long, left: Long, right: Long): Boolean {
+    return x in left..right
+}
+
+fun inCharWithNullableParameter(x: Char?, left: Char, right: Char): Boolean {
+    return x!! in left..right
+}
+
+// 0 INVOKESPECIAL
+// 0 NEW
+// 1 INVOKEVIRTUAL java/lang/Character.charValue
+// 1 INVOKEVIRTUAL
+// 0 CHECKCAST
+// 0 INVOKEINTERFACE
+// 0 <init>

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

@@ -1292,6 +1292,27 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/bytecodeText/ranges")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Ranges extends AbstractBytecodeTextTest {
+        public void testAllFilesPresentInRanges() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/ranges"), Pattern.compile("^(.+)\\.kt$"), true);
+        }
+
+        @TestMetadata("inNonMatchingRange.kt")
+        public void testInNonMatchingRange() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/ranges/inNonMatchingRange.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("inOptimizableRange.kt")
+        public void testInOptimizableRange() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/ranges/inOptimizableRange.kt");
+            doTest(fileName);
+        }
+    }
+
     @TestMetadata("compiler/testData/codegen/bytecodeText/statements")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

core/descriptors/src/org/jetbrains/kotlin/builtins/KotlinBuiltIns.java

@@ -402,8 +402,8 @@ public abstract class KotlinBuiltIns {
     public Set<DeclarationDescriptor> getIntegralRanges() {
         return SetsKt.<DeclarationDescriptor>setOf(
                 getBuiltInClassByName("CharRange", rangesPackageFragment),
-                getBuiltInClassByName("IntRange", rangesPackageFragment)
+                getBuiltInClassByName("IntRange", rangesPackageFragment),
-                // TODO: contains in LongRange should be optimized too
+                getBuiltInClassByName("LongRange", rangesPackageFragment)
         );
     }