src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.nodes/src/org/graalvm/compiler/nodes/extended/SwitchNode.java

*** 26,35 ****
--- 26,36 ----
  import static org.graalvm.compiler.nodeinfo.NodeSize.SIZE_UNKNOWN;
  
  import java.util.Arrays;
  
  import org.graalvm.compiler.core.common.type.AbstractPointerStamp;
+ import org.graalvm.compiler.core.common.type.Stamp;
  import org.graalvm.compiler.core.common.type.StampFactory;
  import org.graalvm.compiler.debug.GraalError;
  import org.graalvm.compiler.graph.Node;
  import org.graalvm.compiler.graph.NodeClass;
  import org.graalvm.compiler.graph.NodeSuccessorList;
*** 88,107 ****
--- 89,146 ----
          assert total > 0.999 && total < 1.001 : "Total " + total;
          return true;
      }
  
      @Override
+     public int getSuccessorCount() {
+         return successors.count();
+     }
+ 
+     @Override
      public double probability(AbstractBeginNode successor) {
          double sum = 0;
          for (int i = 0; i < keySuccessors.length; i++) {
              if (successors.get(keySuccessors[i]) == successor) {
                  sum += keyProbabilities[i];
              }
          }
          return sum;
      }
  
+     @Override
+     public boolean setProbability(AbstractBeginNode successor, double value) {
+         assert value <= 1.0 && value >= 0.0 : value;
+         assert assertProbabilities();
+ 
+         double sum = 0;
+         double otherSum = 0;
+         for (int i = 0; i < keySuccessors.length; i++) {
+             if (successors.get(keySuccessors[i]) == successor) {
+                 sum += keyProbabilities[i];
+             } else {
+                 otherSum += keyProbabilities[i];
+             }
+         }
+ 
+         if (otherSum == 0 || sum == 0) {
+             // Cannot correctly adjust probabilities.
+             return false;
+         }
+ 
+         double delta = value - sum;
+ 
+         for (int i = 0; i < keySuccessors.length; i++) {
+             if (successors.get(keySuccessors[i]) == successor) {
+                 keyProbabilities[i] = Math.max(0.0, keyProbabilities[i] + (delta * keyProbabilities[i]) / sum);
+             } else {
+                 keyProbabilities[i] = Math.max(0.0, keyProbabilities[i] - (delta * keyProbabilities[i]) / otherSum);
+             }
+         }
+         assert assertProbabilities();
+         return true;
+     }
+ 
      public ValueNode value() {
          return value;
      }
  
      public abstract boolean isSorted();
*** 199,204 ****
--- 238,245 ----
              }
          }
          tool.addToWorkList(blockSuccessor(survivingEdge));
          graph().removeSplit(this, blockSuccessor(survivingEdge));
      }
+ 
+     public abstract Stamp getValueStampForSuccessor(AbstractBeginNode beginNode);
  }