Multi-declarations supported in all flavours of for loops

(KT-2635 Support iteration with multiple range variable in JVM back-end) #KT-2635 Fixed
2012-08-31 12:32:07 +04:00
parent f1c4ab7188
commit ce1320bab3
46 changed files with 1603 additions and 74 deletions
@@ -411,23 +411,39 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
    }

    @Override
-    public StackValue visitForExpression(JetForExpression expression, StackValue receiver) {
-        final JetExpression loopRange = expression.getLoopRange();
-        final JetType expressionType = bindingContext.get(BindingContext.EXPRESSION_TYPE, loopRange);
-        assert expressionType != null;
-        Type loopRangeType = asmType(expressionType);
-        if (loopRangeType.getSort() == Type.ARRAY) {
-            generateForLoop(new ForInArrayLoopGenerator(expression));
+    public StackValue visitForExpression(JetForExpression forExpression, StackValue receiver) {
+        // Is it a "1..2" or so
+        RangeCodegenUtil.BinaryCall binaryCall = RangeCodegenUtil.getRangeAsBinaryCall(forExpression);
+        if (binaryCall != null) {
+            ResolvedCall<? extends CallableDescriptor> resolvedCall = bindingContext.get(RESOLVED_CALL, binaryCall.op);
+            assert resolvedCall != null;
+
+            CallableDescriptor rangeTo = resolvedCall.getResultingDescriptor();
+            if (RangeCodegenUtil.isOptimizableRangeTo(rangeTo)
+                // todo: currently, only Int ranges are supported, but all primitives must be optimized
+                && RangeCodegenUtil.isIntRange(rangeTo.getReturnType())
+            ) {
+                generateForLoop(new ForInIntRangeLiteralLoopGenerator(forExpression, binaryCall, resolvedCall));
+                return StackValue.none();
+            }
+        }
+
+        final JetExpression loopRange = forExpression.getLoopRange();
+        final JetType loopRangeType = bindingContext.get(BindingContext.EXPRESSION_TYPE, loopRange);
+        assert loopRangeType != null;
+        Type asmLoopRangeType = asmType(loopRangeType);
+        if (asmLoopRangeType.getSort() == Type.ARRAY) {
+            generateForLoop(new ForInArrayLoopGenerator(forExpression));
            return StackValue.none();
        }
        else {
-            final DeclarationDescriptor descriptor = expressionType.getConstructor().getDeclarationDescriptor();
-            if (isClass(descriptor, "IntRange")) {       // TODO IntRange subclasses (now IntRange is final)
-                new ForInRangeLoopGenerator(expression, loopRangeType).invoke();
+            // todo: Only IntRange optimized so far
+            if (RangeCodegenUtil.isIntRange(loopRangeType)) {
+                generateForLoop(new ForInIntRangeInstanceLoopGenerator(forExpression));
                return StackValue.none();
            }

-            generateForLoop(new IteratorForLoopGenerator(expression));
+            generateForLoop(new IteratorForLoopGenerator(forExpression));
            return StackValue.none();
        }
    }
@@ -780,84 +796,135 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
        }
    }

-    private class ForInRangeLoopGenerator extends IntrinsicForLoopGenerator {
-        private int myCountVar;
-        private int myDeltaVar;
-        private int myIndexVar;
+    private class ForInIntRangeLiteralLoopGenerator extends AbstractForLoopGenerator {
+        private final RangeCodegenUtil.BinaryCall rangeCall;
+        private final ResolvedCall<? extends CallableDescriptor> resolvedCall;
+        private Type asmElementType;
+        private int indexVar;
+        private int lastVar;

-        public ForInRangeLoopGenerator(JetForExpression expression, Type loopRangeType) {
-            super(expression, loopRangeType);
+        private ForInIntRangeLiteralLoopGenerator(
+                @NotNull JetForExpression forExpression,
+                @NotNull RangeCodegenUtil.BinaryCall rangeCall,
+                ResolvedCall<? extends CallableDescriptor> resolvedCall
+        ) {
+            super(forExpression);
+            this.rangeCall = rangeCall;
+            this.resolvedCall = resolvedCall;
        }

        @Override
-        protected void generatePrologue() {
-            myIndexVar = lookupLocalIndex(parameterDescriptor);
-            myCountVar = myFrameMap.enterTemp(Type.INT_TYPE);
-            myDeltaVar = myFrameMap.enterTemp(Type.INT_TYPE);
-            if (isIntRangeExpr(expression.getLoopRange())) {
-                JetBinaryExpression rangeExpression = (JetBinaryExpression) expression.getLoopRange();
-                //noinspection ConstantConditions
-                gen(rangeExpression.getLeft(), Type.INT_TYPE);
-                v.store(myIndexVar, Type.INT_TYPE);
-                gen(rangeExpression.getRight(), Type.INT_TYPE);
-                v.store(myCountVar, Type.INT_TYPE);
+        public void beforeLoop() {
+            JetType elementType = RangeCodegenUtil.getPrimitiveRangeElementType(resolvedCall.getResultingDescriptor().getReturnType());
+            assert elementType != null;
+            asmElementType = asmType(elementType);

-                v.load(myCountVar, Type.INT_TYPE);
-                v.load(myIndexVar, Type.INT_TYPE);
-                v.sub(Type.INT_TYPE);
-                v.iconst(1);
-                v.add(Type.INT_TYPE);
-                v.store(myCountVar, Type.INT_TYPE);
+            indexVar = myFrameMap.enterTemp(asmElementType);
+            gen(rangeCall.left, asmElementType);
+            v.store(indexVar, asmElementType);

-                v.load(myCountVar, Type.INT_TYPE);
-                v.iflt(end);
-
-                v.iconst(1);
-                v.store(myDeltaVar, Type.INT_TYPE);
-            }
-            else {
-                gen(expression.getLoopRange(), loopRangeType);
-                v.dup();
-                v.dup();
-
-                v.invokevirtual("jet/IntRange", "getStart", "()I");
-                v.store(myIndexVar, Type.INT_TYPE);
-                v.invokevirtual("jet/IntRange", "getSize", "()I");
-                v.store(myCountVar, Type.INT_TYPE);
-
-                v.invokevirtual("jet/IntRange", "getIsReversed", "()Z");
-                Label down = new Label();
-
-                v.ifne(down);
-                v.iconst(1);
-                Label initEnd = new Label();
-                v.goTo(initEnd);
-                v.mark(down);
-                v.iconst(-1);
-                v.mark(initEnd);
-                v.store(myDeltaVar, Type.INT_TYPE);
-            }
+            lastVar = myFrameMap.enterTemp(asmElementType);
+            gen(rangeCall.right, asmElementType);
+            v.store(lastVar, asmElementType);
        }

        @Override
-        protected void generateCondition(Type asmParamType, Label end) {
-            v.load(myCountVar, Type.INT_TYPE);
-            v.ifeq(end);
+        public void conditionAndJump(@NotNull Label loopExit) {
+            v.load(indexVar, asmElementType);
+            v.load(lastVar, asmElementType);
+            v.ificmpgt(loopExit);
        }

        @Override
-        protected void generateIncrement() {
-            v.load(myIndexVar, Type.INT_TYPE);
-            v.load(myDeltaVar, Type.INT_TYPE);
+        protected void assignToLoopParameter(int parameterIndex) {
+            // todo: don't create a temp variable if this is not a multi-decl for
+            v.load(indexVar, asmElementType);
+            v.store(parameterIndex, asmElementType);
+        }
+
+        @Override
+        public void afterBody() {
+            v.iinc(indexVar, 1);
+            super.afterBody();
+        }
+
+        @Override
+        public void afterLoop() {
+            myFrameMap.leaveTemp(asmElementType); // lastVar
+            myFrameMap.leaveTemp(asmElementType); // indexVar
+        }
+    }
+
+    private class ForInIntRangeInstanceLoopGenerator extends AbstractForLoopGenerator {
+        private int indexVar;
+        private int countVar;
+        private int deltaVar;
+
+        private ForInIntRangeInstanceLoopGenerator(
+                @NotNull JetForExpression forExpression
+        ) {
+            super(forExpression);
+        }
+
+        @Override
+        public void beforeLoop() {
+            JetType loopRangeType = bindingContext.get(EXPRESSION_TYPE, forExpression.getLoopRange());
+            assert loopRangeType != null;
+            Type asmLoopRangeType = asmType(loopRangeType);
+            gen(forExpression.getLoopRange(), asmLoopRangeType);
+            v.dup();
+            v.dup();
+
+            indexVar = myFrameMap.enterTemp(Type.INT_TYPE);
+            v.invokevirtual(JET_INT_RANGE_TYPE.getInternalName(), "getStart", "()I");
+            v.store(indexVar, Type.INT_TYPE);
+
+            countVar = myFrameMap.enterTemp(Type.INT_TYPE);
+            v.invokevirtual(JET_INT_RANGE_TYPE.getInternalName(), "getSize", "()I");
+            v.store(countVar, Type.INT_TYPE);
+
+            deltaVar = myFrameMap.enterTemp(Type.INT_TYPE);
+            v.invokevirtual(JET_INT_RANGE_TYPE.getInternalName(), "getIsReversed", "()Z");
+            Label down = new Label();
+
+            v.ifne(down);
+            v.iconst(1);
+            Label initEnd = new Label();
+            v.goTo(initEnd);
+            v.mark(down);
+            v.iconst(-1);
+            v.mark(initEnd);
+            v.store(deltaVar, Type.INT_TYPE);
+        }
+
+        @Override
+        public void conditionAndJump(@NotNull Label loopExit) {
+            v.load(countVar, Type.INT_TYPE);
+            v.ifeq(loopExit);
+        }
+
+        @Override
+        protected void assignToLoopParameter(int parameterIndex) {
+            // todo: no temp var when this is not a multi-declaration for-loop
+            v.load(indexVar, Type.INT_TYPE);
+            v.store(parameterIndex, Type.INT_TYPE);
+        }
+
+        @Override
+        public void afterBody() {
+            v.load(indexVar, Type.INT_TYPE);
+            v.load(deltaVar, Type.INT_TYPE);
            v.add(Type.INT_TYPE);
-            v.store(myIndexVar, Type.INT_TYPE);
-            v.iinc(myCountVar, -1);
+            v.store(indexVar, Type.INT_TYPE);
+            v.iinc(countVar, -1);
+            super.afterBody();
        }

        @Override
-        protected void cleanupTemp() {
-            myFrameMap.leaveTemp(Type.INT_TYPE);
-            myFrameMap.leaveTemp(Type.INT_TYPE);
+        public void afterLoop() {
+            myFrameMap.leaveTemp(Type.INT_TYPE); // deltaVar
+            myFrameMap.leaveTemp(Type.INT_TYPE); // countVar
+            myFrameMap.leaveTemp(Type.INT_TYPE); // indexVar
        }
    }

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2010-2012 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.codegen;
+
+import com.google.common.collect.ImmutableMap;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.descriptors.CallableDescriptor;
+import org.jetbrains.jet.lang.descriptors.ClassifierDescriptor;
+import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.types.JetType;
+import org.jetbrains.jet.lang.types.lang.JetStandardLibrary;
+
+import java.util.List;
+
+/**
+ * @author abreslav
+ */
+public class RangeCodegenUtil {
+    private static final ImmutableMap<String, JetType> RANGE_TO_ELEMENT_TYPE = ImmutableMap.<String, JetType>builder()
+            .put("ByteRange", JetStandardLibrary.getInstance().getByteType())
+            .put("ShortRange", JetStandardLibrary.getInstance().getShortType())
+            .put("IntRange", JetStandardLibrary.getInstance().getIntType())
+            .put("LongRange", JetStandardLibrary.getInstance().getLongType())
+            .put("FloatRange", JetStandardLibrary.getInstance().getFloatType())
+            .put("DoubleRange", JetStandardLibrary.getInstance().getDoubleType())
+            .put("CharRange", JetStandardLibrary.getInstance().getCharType())
+            .build();
+
+    private RangeCodegenUtil() {}
+
+    public static boolean isIntRange(JetType rangeType) {
+        return !rangeType.isNullable()
+               && JetStandardLibrary.getInstance().getIntType().equals(getPrimitiveRangeElementType(rangeType));
+    }
+
+    @Nullable
+    public static BinaryCall getRangeAsBinaryCall(@NotNull JetForExpression forExpression) {
+        JetExpression rangeExpression = forExpression.getLoopRange();
+        assert rangeExpression != null;
+        JetExpression loopRange = JetPsiUtil.deparenthesize(rangeExpression);
+        if (loopRange instanceof JetQualifiedExpression) {
+            // a.rangeTo(b)
+            JetQualifiedExpression qualifiedExpression = (JetQualifiedExpression) loopRange;
+            JetExpression selector = qualifiedExpression.getSelectorExpression();
+            if (selector instanceof JetCallExpression) {
+                JetCallExpression callExpression = (JetCallExpression) selector;
+                List<? extends ValueArgument> arguments = callExpression.getValueArguments();
+                if (arguments.size() == 1) {
+                    return new BinaryCall(qualifiedExpression.getReceiverExpression(), callExpression.getCalleeExpression(),
+                                          arguments.get(0).getArgumentExpression());
+                }
+            }
+        }
+        else if (loopRange instanceof JetBinaryExpression) {
+            // a rangeTo b
+            // a .. b
+            JetBinaryExpression binaryExpression = (JetBinaryExpression) loopRange;
+            return new BinaryCall(binaryExpression.getLeft(), binaryExpression.getOperationReference(), binaryExpression.getRight());
+
+        }
+        return null;
+    }
+
+    @Nullable
+    public static JetType getPrimitiveRangeElementType(JetType rangeType) {
+        ClassifierDescriptor declarationDescriptor = rangeType.getConstructor().getDeclarationDescriptor();
+        assert declarationDescriptor != null;
+        if (declarationDescriptor != JetStandardLibrary.getInstance().getLibraryScope().getClassifier(declarationDescriptor.getName())) {
+            // Must be a standard library class
+            return null;
+        }
+        return RANGE_TO_ELEMENT_TYPE.get(declarationDescriptor.getName().getName());
+    }
+
+    public static boolean isOptimizableRangeTo(CallableDescriptor rangeTo) {
+        if ("rangeTo".equals(rangeTo.getName().getName())) {
+            if (CodegenUtil.isPrimitiveNumberClassDescriptor(rangeTo.getContainingDeclaration())) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    public static class BinaryCall {
+        public final JetExpression left;
+        public final JetExpression op;
+        public final JetExpression right;
+
+        private BinaryCall(JetExpression left, JetExpression op, JetExpression right) {
+            this.left = left;
+            this.op = op;
+            this.right = right;
+        }
+    }
+}