KT-933 Optimized iteration over progressions.

#KT-933 fixed
2013-02-11 22:27:02 +04:00
parent 14f3144533
commit 64c89836fd
2 changed files with 159 additions and 0 deletions
@@ -56,6 +56,7 @@ import org.jetbrains.jet.lang.resolve.java.AsmTypeConstants;
 import org.jetbrains.jet.lang.resolve.java.JvmAbi;
 import org.jetbrains.jet.lang.resolve.java.JvmClassName;
 import org.jetbrains.jet.lang.resolve.java.JvmPrimitiveType;
+import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.receivers.*;
 import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.lang.types.checker.JetTypeChecker;
@@ -484,6 +485,11 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
                return StackValue.none();
            }

+            if (RangeCodegenUtil.isProgression(loopRangeType)) {
+                generateForLoop(new ForInProgressionExpressionLoopGenerator(forExpression));
+                return StackValue.none();
+            }
+
            generateForLoop(new IteratorForLoopGenerator(forExpression));
            return StackValue.none();
        }
@@ -939,6 +945,140 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
        }
    }

+    private class ForInProgressionExpressionLoopGenerator extends AbstractForLoopGenerator {
+        private int endVar;
+        private int incrementVar;
+        private Type incrementType;
+
+        private ForInProgressionExpressionLoopGenerator(@NotNull JetForExpression forExpression) {
+            super(forExpression);
+        }
+
+        @Override
+        public void beforeLoop() {
+            super.beforeLoop();
+
+            endVar = createLoopTempVariable(asmElementType);
+            incrementVar = createLoopTempVariable(asmElementType);
+
+            JetType loopRangeType = bindingContext.get(EXPRESSION_TYPE, forExpression.getLoopRange());
+            assert loopRangeType != null;
+            Type asmLoopRangeType = asmType(loopRangeType);
+
+            Collection<VariableDescriptor> incrementProp = loopRangeType.getMemberScope().getProperties(Name.identifier("increment"));
+            assert incrementProp.size() == 1 : loopRangeType + " " + incrementProp.size();
+            incrementType = asmType(incrementProp.iterator().next().getType());
+
+            gen(forExpression.getLoopRange(), asmLoopRangeType);
+            v.dup();
+            v.dup();
+
+            generateRangeOrProgressionProperty(asmLoopRangeType, "getStart", asmElementType, loopParameterVar);
+            generateRangeOrProgressionProperty(asmLoopRangeType, "getEnd", asmElementType, endVar);
+            generateRangeOrProgressionProperty(asmLoopRangeType, "getIncrement", incrementType, incrementVar);
+        }
+
+        @Override
+        public void conditionAndJump(@NotNull Label loopExit) {
+            v.load(loopParameterVar, asmElementType);
+            v.load(incrementVar, asmElementType);
+
+            Label negativeIncrement = new Label();
+            Label afterIf = new Label();
+
+            int sort = asmElementType.getSort();
+            switch (sort) {
+                case Type.INT:
+                case Type.CHAR:
+                case Type.BYTE:
+                case Type.SHORT:
+                    v.ifle(negativeIncrement); // if increment < 0, jump
+
+                    // increment > 0
+                    v.load(endVar, asmElementType);
+                    v.ificmpgt(loopExit);
+                    v.goTo(afterIf);
+
+                    // increment < 0
+                    v.visitLabel(negativeIncrement);
+                    v.load(endVar, asmElementType);
+                    v.ificmplt(loopExit);
+                    v.visitLabel(afterIf);
+
+                    break;
+
+                case Type.LONG:
+                    v.lconst(0L);
+                    v.lcmp();
+                    v.ifle(negativeIncrement); // if increment < 0, jump
+
+                    // increment > 0
+                    v.load(endVar, asmElementType);
+                    v.lcmp();
+                    v.ifgt(loopExit);
+                    v.goTo(afterIf);
+
+                    // increment < 0
+                    v.visitLabel(negativeIncrement);
+                    v.load(endVar, asmElementType);
+                    v.lcmp();
+                    v.iflt(loopExit);
+                    v.visitLabel(afterIf);
+
+                    break;
+
+                case Type.DOUBLE:
+                case Type.FLOAT:
+                    if (sort == Type.DOUBLE) {
+                        v.dconst(0.0);
+                    }
+                    else {
+                        v.fconst(0.0f);
+                    }
+                    v.cmpl(incrementType);
+                    v.ifle(negativeIncrement); // if increment < 0, jump
+
+                    // increment > 0
+                    v.load(endVar, asmElementType);
+                    v.cmpg(asmElementType); // if loop parameter is NaN, exit from loop, as well
+                    v.ifgt(loopExit);
+                    v.goTo(afterIf);
+
+                    // increment < 0
+                    v.visitLabel(negativeIncrement);
+                    v.load(endVar, asmElementType);
+                    v.cmpl(asmElementType); // if loop parameter is NaN, exit from loop, as well
+                    v.iflt(loopExit);
+                    v.visitLabel(afterIf);
+
+                    break;
+
+                default:
+                    throw new IllegalStateException("Unexpected range element type: " + asmElementType);
+            }
+
+        }
+
+        @Override
+        protected void assignToLoopParameter() {
+        }
+
+        @Override
+        protected void increment() {
+
+            v.load(loopParameterVar, asmElementType);
+            v.load(incrementVar, asmElementType);
+            v.add(asmElementType);
+
+            int sort = asmElementType.getSort();
+            if (sort == Type.CHAR || sort == Type.BYTE || sort == Type.SHORT) {
+                StackValue.coerce(Type.INT_TYPE, asmElementType, v);
+            }
+
+            v.store(loopParameterVar, asmElementType);
+        }
+    }
+
    @Override
    public StackValue visitBreakExpression(JetBreakExpression expression, StackValue receiver) {
        JetSimpleNameExpression labelElement = expression.getTargetLabel();
@@ -34,13 +34,17 @@ import static org.jetbrains.jet.codegen.AsmUtil.isPrimitiveNumberClassDescriptor

 public class RangeCodegenUtil {
    private static final ImmutableMap<FqName, PrimitiveType> RANGE_TO_ELEMENT_TYPE;
+    private static final ImmutableMap<FqName, PrimitiveType> PROGRESSION_TO_ELEMENT_TYPE;

    static {
        ImmutableMap.Builder<FqName, PrimitiveType> rangeBuilder = ImmutableMap.builder();
+        ImmutableMap.Builder<FqName, PrimitiveType> progressionBuilder = ImmutableMap.builder();
        for (PrimitiveType primitiveType : PrimitiveType.values()) {
            rangeBuilder.put(primitiveType.getRangeClassName(), primitiveType);
+            progressionBuilder.put(primitiveType.getProgressionClassName(), primitiveType);
        }
        RANGE_TO_ELEMENT_TYPE = rangeBuilder.build();
+        PROGRESSION_TO_ELEMENT_TYPE = progressionBuilder.build();
    }

    private RangeCodegenUtil() {}
@@ -49,6 +53,10 @@ public class RangeCodegenUtil {
        return !rangeType.isNullable() && getPrimitiveRangeElementType(rangeType) != null;
    }

+    public static boolean isProgression(JetType rangeType) {
+        return !rangeType.isNullable() && getPrimitiveProgressionElementType(rangeType) != null;
+    }
+
    @Nullable
    public static BinaryCall getRangeAsBinaryCall(@NotNull JetForExpression forExpression) {
        // We are looking for rangeTo() calls
@@ -90,6 +98,17 @@ public class RangeCodegenUtil {
        return RANGE_TO_ELEMENT_TYPE.get(DescriptorUtils.getFQName(declarationDescriptor).toSafe());
    }

+    @Nullable
+    private static PrimitiveType getPrimitiveProgressionElementType(JetType rangeType) {
+        ClassifierDescriptor declarationDescriptor = rangeType.getConstructor().getDeclarationDescriptor();
+        assert declarationDescriptor != null;
+        if (declarationDescriptor != KotlinBuiltIns.getInstance().getBuiltInsScope().getClassifier(declarationDescriptor.getName())) {
+            // Must be a standard library class
+            return null;
+        }
+        return PROGRESSION_TO_ELEMENT_TYPE.get(DescriptorUtils.getFQName(declarationDescriptor).toSafe());
+    }
+
    public static boolean isOptimizableRangeTo(CallableDescriptor rangeTo) {
        if ("rangeTo".equals(rangeTo.getName().getName())) {
            if (isPrimitiveNumberClassDescriptor(rangeTo.getContainingDeclaration())) {