Jump on 'continue' to condition entry points for all loops

compiler/frontend/src/org/jetbrains/jet/lang/cfg/BreakableBlockInfo.java

@@ -16,17 +16,22 @@
 
 package org.jetbrains.jet.lang.cfg;
 
+import com.google.common.collect.Sets;
 import org.jetbrains.jet.lang.psi.JetElement;
 
+import java.util.Set;
+
 public class BreakableBlockInfo extends BlockInfo {
     private final JetElement element;
     private final Label entryPoint;
     private final Label exitPoint;
+    protected final Set<Label> referablePoints;
 
     public BreakableBlockInfo(JetElement element, Label entryPoint, Label exitPoint) {
         this.element = element;
         this.entryPoint = entryPoint;
         this.exitPoint = exitPoint;
+        referablePoints = Sets.newHashSet(entryPoint, exitPoint);
     }
 
     public JetElement getElement() {
@@ -40,4 +45,8 @@ public class BreakableBlockInfo extends BlockInfo {
     public Label getExitPoint() {
         return exitPoint;
     }
+
+    public Set<Label> getReferablePoints() {
+        return referablePoints;
+    }
 }

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowProcessor.java

@@ -953,11 +953,7 @@ public class JetControlFlowProcessor {
             JetElement loop = getCorrespondingLoop(expression);
             if (loop != null) {
                 checkJumpDoesNotCrossFunctionBoundary(expression, loop);
-                if (loop instanceof JetDoWhileExpression) {
-                    builder.jump(builder.getConditionEntryPoint(loop), expression);
-                } else {
-                    builder.jump(builder.getEntryPoint(loop), expression);
-                }
+                builder.jump(builder.getConditionEntryPoint(loop), expression);
             }
         }
 

compiler/frontend/src/org/jetbrains/jet/lang/cfg/LoopInfo.java

@@ -16,8 +16,11 @@
 
 package org.jetbrains.jet.lang.cfg;
 
+import com.google.common.collect.Sets;
 import org.jetbrains.jet.lang.psi.JetElement;
 
+import java.util.Set;
+
 public class LoopInfo extends BreakableBlockInfo {
     private final Label bodyEntryPoint;
     private final Label conditionEntryPoint;
@@ -26,6 +29,8 @@ public class LoopInfo extends BreakableBlockInfo {
         super(element, entryPoint, exitPoint);
         this.bodyEntryPoint = bodyEntryPoint;
         this.conditionEntryPoint = conditionEntryPoint;
+        referablePoints.add(bodyEntryPoint);
+        referablePoints.add(conditionEntryPoint);
     }
 
     public Label getBodyEntryPoint() {

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/JetControlFlowInstructionsGenerator.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.jet.lang.cfg.pseudocode;
 
+import com.google.common.collect.Sets;
 import com.intellij.util.containers.Stack;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -246,8 +247,7 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
                 BlockInfo blockInfo = allBlocks.get(i);
                 if (blockInfo instanceof BreakableBlockInfo) {
                     BreakableBlockInfo breakableBlockInfo = (BreakableBlockInfo) blockInfo;
-                    if (jumpTarget == breakableBlockInfo.getExitPoint() || jumpTarget == breakableBlockInfo.getEntryPoint()
-                        || jumpTarget == error) {
+                    if (breakableBlockInfo.getReferablePoints().contains(jumpTarget) || jumpTarget == error) {
                         for (int j = finallyBlocks.size() - 1; j >= 0; j--) {
                             finallyBlocks.get(j).generateFinallyBlock();
                         }