# HG changeset patch
# User Gilles Duboscq <duboscq@ssw.jku.at>
# Date 1397755837 -7200
# Node ID f52e47b757b92f5a84f2b26c8726ae8cf353e532
# Parent  349a0371b3b6b837f4cac4a28dca13c5f2efc80b
Loop transform: mark states from all NodeWithStates not just StateSplit

diff -r 349a0371b3b6 -r f52e47b757b9 graal/com.oracle.graal.loop/src/com/oracle/graal/loop/LoopFragment.java
--- a/graal/com.oracle.graal.loop/src/com/oracle/graal/loop/LoopFragment.java	Thu Apr 17 19:29:54 2014 +0200
+++ b/graal/com.oracle.graal.loop/src/com/oracle/graal/loop/LoopFragment.java	Thu Apr 17 19:30:37 2014 +0200
@@ -31,6 +31,7 @@
 import com.oracle.graal.nodes.*;
 import com.oracle.graal.nodes.cfg.*;
 import com.oracle.graal.nodes.java.*;
+import com.oracle.graal.nodes.spi.*;
 import com.oracle.graal.nodes.virtual.*;
 
 public abstract class LoopFragment {
@@ -164,11 +165,9 @@
                 if (n instanceof Invoke) {
                     nodes.mark(((Invoke) n).callTarget());
                 }
-                if (n instanceof StateSplit) {
-                    FrameState stateAfter = ((StateSplit) n).stateAfter();
-                    if (stateAfter != null) {
-                        stateAfter.applyToVirtual(node -> nodes.mark(node));
-                    }
+                if (n instanceof NodeWithState) {
+                    NodeWithState withState = (NodeWithState) n;
+                    withState.states().forEach(state -> state.applyToVirtual(node -> nodes.mark(node)));
                 }
                 nodes.mark(n);
             }
@@ -334,7 +333,7 @@
                  * VirtualState nodes contained in the old exit's state may be shared by other
                  * dominated VirtualStates. Those dominated virtual states need to see the
                  * proxy->phi update that are applied below.
-                 * 
+                 *
                  * We now update the original fragment's nodes accordingly:
                  */
                 state.applyToVirtual(node -> original.nodes.clear(node));