Useful method extracted + visiting strategy added

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

@@ -18,15 +18,13 @@ package org.jetbrains.jet.lang.cfg;
 
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Queues;
 import com.google.common.collect.Sets;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.cfg.pseudocode.*;
 
-import java.util.Collection;
+import java.util.*;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
 
 import static org.jetbrains.jet.lang.cfg.PseudocodeTraverser.TraversalOrder.FORWARD;
 
@@ -236,4 +234,37 @@ public class PseudocodeTraverser {
             return result;
         }
     }
+
+    public interface InstructionHandler {
+        // true to continue traversal
+        boolean handle(@NotNull Instruction instruction);
+    }
+
+    // returns false when interrupted by handler
+    public static boolean traverseFollowingInstructions(
+            @NotNull Instruction rootInstruction,
+            @NotNull Set<Instruction> visited,
+            @NotNull TraversalOrder order,
+            @Nullable InstructionHandler handler
+    ) {
+        Deque<Instruction> stack = Queues.newArrayDeque();
+        stack.push(rootInstruction);
+
+        while (!stack.isEmpty()) {
+            Instruction instruction = stack.pop();
+            visited.add(instruction);
+
+            Collection<Instruction> followingInstructions =
+                    order == FORWARD ? instruction.getNextInstructions() : instruction.getPreviousInstructions();
+
+            for (Instruction followingInstruction : followingInstructions) {
+                if (followingInstruction != null && !visited.contains(followingInstruction)) {
+                    if (handler != null && !handler.handle(instruction)) return false;
+                    stack.push(followingInstruction);
+                }
+            }
+        }
+        return true;
+    }
+
 }

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

@@ -21,11 +21,15 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.cfg.Label;
 import org.jetbrains.jet.lang.cfg.LoopInfo;
+import org.jetbrains.jet.lang.cfg.PseudocodeTraverser;
 import org.jetbrains.jet.lang.psi.JetElement;
 import org.jetbrains.jet.lang.psi.JetExpression;
 
 import java.util.*;
 
+import static org.jetbrains.jet.lang.cfg.PseudocodeTraverser.TraversalOrder.BACKWARD;
+import static org.jetbrains.jet.lang.cfg.PseudocodeTraverser.TraversalOrder.FORWARD;
+
 public class PseudocodeImpl implements Pseudocode {
 
     public class PseudocodeLabel implements Label {
@@ -134,13 +138,13 @@ public class PseudocodeImpl implements Pseudocode {
     @Override
     public List<Instruction> getReversedInstructions() {
         LinkedHashSet<Instruction> traversedInstructions = Sets.newLinkedHashSet();
-        traverseFollowingInstructions(sinkInstruction, traversedInstructions, false);
+        PseudocodeTraverser.traverseFollowingInstructions(sinkInstruction, traversedInstructions, BACKWARD, null);
         if (traversedInstructions.size() < instructions.size()) {
             List<Instruction> simplyReversedInstructions = Lists.newArrayList(instructions);
             Collections.reverse(simplyReversedInstructions);
             for (Instruction instruction : simplyReversedInstructions) {
                 if (!traversedInstructions.contains(instruction)) {
-                    traverseFollowingInstructions(instruction, traversedInstructions, false);
+                    PseudocodeTraverser.traverseFollowingInstructions(instruction, traversedInstructions, BACKWARD, null);
                 }
             }
         }
@@ -311,7 +315,7 @@ public class PseudocodeImpl implements Pseudocode {
 
     private Set<Instruction> collectReachableInstructions() {
         Set<Instruction> visited = Sets.newHashSet();
-        traverseFollowingInstructions(getEnterInstruction(), visited, true);
+        PseudocodeTraverser.traverseFollowingInstructions(getEnterInstruction(), visited, FORWARD, null);
         if (!visited.contains(getExitInstruction())) {
             visited.add(getExitInstruction());
         }
@@ -324,29 +328,6 @@ public class PseudocodeImpl implements Pseudocode {
         return visited;
     }
 
-    private static void traverseFollowingInstructions(
-            @NotNull Instruction rootInstruction,
-            @NotNull Set<Instruction> visited,
-            boolean directOrder
-    ) {
-        Deque<Instruction> stack = Queues.newArrayDeque();
-        stack.push(rootInstruction);
-
-        while (!stack.isEmpty()) {
-            Instruction instruction = stack.pop();
-            visited.add(instruction);
-
-            Collection<Instruction> followingInstructions =
-                    directOrder ? instruction.getNextInstructions() : instruction.getPreviousInstructions();
-
-            for (Instruction followingInstruction : followingInstructions) {
-                if (followingInstruction != null && !visited.contains(followingInstruction)) {
-                    stack.push(followingInstruction);
-                }
-            }
-        }
-    }
-
     private void markDeadInstructions() {
         Set<Instruction> instructionSet = Sets.newHashSet(instructions);
         for (Instruction instruction : mutableInstructionList) {