package org.graalvm.compiler.nodes.extended;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.graalvm.compiler.core.common.spi.ConstantFieldProvider;
import org.graalvm.compiler.core.common.type.IntegerStamp;
import org.graalvm.compiler.core.common.type.PrimitiveStamp;
import org.graalvm.compiler.core.common.type.Stamp;
import org.graalvm.compiler.core.common.type.StampFactory;
import org.graalvm.compiler.graph.NodeClass;
import org.graalvm.compiler.graph.spi.Simplifiable;
import org.graalvm.compiler.graph.spi.SimplifierTool;
import org.graalvm.compiler.nodeinfo.NodeInfo;
import org.graalvm.compiler.nodes.AbstractBeginNode;
import org.graalvm.compiler.nodes.ConstantNode;
import org.graalvm.compiler.nodes.FixedGuardNode;
import org.graalvm.compiler.nodes.FixedNode;
import org.graalvm.compiler.nodes.FixedWithNextNode;
import org.graalvm.compiler.nodes.LogicNode;
import org.graalvm.compiler.nodes.NodeView;
import org.graalvm.compiler.nodes.ValueNode;
import org.graalvm.compiler.nodes.calc.IntegerBelowNode;
import org.graalvm.compiler.nodes.java.LoadIndexedNode;
import org.graalvm.compiler.nodes.spi.LIRLowerable;
import org.graalvm.compiler.nodes.spi.NodeLIRBuilderTool;
import org.graalvm.compiler.nodes.util.GraphUtil;
import jdk.vm.ci.meta.DeoptimizationAction;
import jdk.vm.ci.meta.DeoptimizationReason;
import jdk.vm.ci.meta.JavaConstant;
import jdk.vm.ci.meta.JavaKind;
@NodeInfo
public final class IntegerSwitchNode extends SwitchNode implements LIRLowerable, Simplifiable {
public static final NodeClass<IntegerSwitchNode> TYPE = NodeClass.create(IntegerSwitchNode.class);
protected final int[] keys;
public IntegerSwitchNode(ValueNode value, AbstractBeginNode[] successors, int[] keys, double[] keyProbabilities, int[] keySuccessors) {
super(TYPE, value, successors, keySuccessors, keyProbabilities);
assert keySuccessors.length == keys.length + 1;
assert keySuccessors.length == keyProbabilities.length;
this.keys = keys;
assert value.stamp(NodeView.DEFAULT) instanceof PrimitiveStamp && value.stamp(NodeView.DEFAULT).getStackKind().isNumericInteger();
assert assertSorted();
}
private boolean assertSorted() {
for (int i = 1; i < keys.length; i++) {
assert keys[i - 1] < keys[i];
}
return true;
}
public IntegerSwitchNode(ValueNode value, int successorCount, int[] keys, double[] keyProbabilities, int[] keySuccessors) {
this(value, new AbstractBeginNode[successorCount], keys, keyProbabilities, keySuccessors);
}
@Override
public boolean isSorted() {
return true;
}
@Override
public JavaConstant keyAt(int i) {
return JavaConstant.forInt(keys[i]);
}
@Override
public int keyCount() {
return keys.length;
}
@Override
public boolean equalKeys(SwitchNode switchNode) {
if (!(switchNode instanceof IntegerSwitchNode)) {
return false;
}
IntegerSwitchNode other = (IntegerSwitchNode) switchNode;
return Arrays.equals(keys, other.keys);
}
@Override
public void generate(NodeLIRBuilderTool gen) {
gen.emitSwitch(this);
}
public AbstractBeginNode successorAtKey(int key) {
return blockSuccessor(successorIndexAtKey(key));
}
public int successorIndexAtKey(int key) {
for (int i = 0; i < keyCount(); i++) {
if (keys[i] == key) {
return keySuccessorIndex(i);
}
}
return keySuccessorIndex(keyCount());
}
@Override
public void simplify(SimplifierTool tool) {
NodeView view = NodeView.from(tool);
if (blockSuccessorCount() == 1) {
tool.addToWorkList(defaultSuccessor());
graph().removeSplitPropagate(this, defaultSuccessor());
} else if (value() instanceof ConstantNode) {
killOtherSuccessors(tool, successorIndexAtKey(value().asJavaConstant().asInt()));
} else if (tryOptimizeEnumSwitch(tool)) {
return;
} else if (tryRemoveUnreachableKeys(tool, value().stamp(view))) {
return;
}
}
static final class KeyData {
final int key;
final double keyProbability;
final int keySuccessor;
KeyData(int key, double keyProbability, int keySuccessor) {
this.key = key;
this.keyProbability = keyProbability;
this.keySuccessor = keySuccessor;
}
}
public boolean tryRemoveUnreachableKeys(SimplifierTool tool, Stamp valueStamp) {
if (!(valueStamp instanceof IntegerStamp)) {
return false;
}
IntegerStamp integerStamp = (IntegerStamp) valueStamp;
if (integerStamp.isUnrestricted()) {
return false;
}
List<KeyData> newKeyDatas = new ArrayList<>(keys.length);
ArrayList<AbstractBeginNode> newSuccessors = new ArrayList<>(blockSuccessorCount());
for (int i = 0; i < keys.length; i++) {
if (integerStamp.contains(keys[i]) && keySuccessor(i) != defaultSuccessor()) {
newKeyDatas.add(new KeyData(keys[i], keyProbabilities[i], addNewSuccessor(keySuccessor(i), newSuccessors)));
}
}
if (newKeyDatas.size() == keys.length) {
return false;
} else if (newKeyDatas.size() == 0) {
if (tool != null) {
tool.addToWorkList(defaultSuccessor());
}
graph().removeSplitPropagate(this, defaultSuccessor());
return true;
} else {
int newDefaultSuccessor = addNewSuccessor(defaultSuccessor(), newSuccessors);
double newDefaultProbability = keyProbabilities[keyProbabilities.length - 1];
doReplace(value(), newKeyDatas, newSuccessors, newDefaultSuccessor, newDefaultProbability);
return true;
}
}
private boolean tryOptimizeEnumSwitch(SimplifierTool tool) {
if (!(value() instanceof LoadIndexedNode)) {
return false;
}
LoadIndexedNode loadIndexed = (LoadIndexedNode) value();
if (loadIndexed.usages().count() > 1) {
return false;
}
assert loadIndexed.usages().first() == this;
ValueNode newValue = loadIndexed.index();
JavaConstant arrayConstant = loadIndexed.array().asJavaConstant();
if (arrayConstant == null || ((ConstantNode) loadIndexed.array()).getStableDimension() != 1 || !((ConstantNode) loadIndexed.array()).isDefaultStable()) {
return false;
}
Integer optionalArrayLength = tool.getConstantReflection().readArrayLength(arrayConstant);
if (optionalArrayLength == null) {
return false;
}
int arrayLength = optionalArrayLength;
Map<Integer, List<Integer>> reverseArrayMapping = new HashMap<>();
for (int i = 0; i < arrayLength; i++) {
JavaConstant elementConstant = tool.getConstantReflection().readArrayElement(arrayConstant, i);
if (elementConstant == null || elementConstant.getJavaKind() != JavaKind.Int) {
return false;
}
int element = elementConstant.asInt();
reverseArrayMapping.computeIfAbsent(element, e -> new ArrayList<>()).add(i);
}
List<KeyData> newKeyDatas = new ArrayList<>(arrayLength);
ArrayList<AbstractBeginNode> newSuccessors = new ArrayList<>(blockSuccessorCount());
for (int i = 0; i < keys.length; i++) {
List<Integer> newKeys = reverseArrayMapping.get(keys[i]);
if (newKeys == null || newKeys.size() == 0) {
continue;
}
double newKeyProbability = keyProbabilities[i] / newKeys.size();
int newKeySuccessor = addNewSuccessor(keySuccessor(i), newSuccessors);
for (int newKey : newKeys) {
newKeyDatas.add(new KeyData(newKey, newKeyProbability, newKeySuccessor));
}
}
int newDefaultSuccessor = addNewSuccessor(defaultSuccessor(), newSuccessors);
double newDefaultProbability = keyProbabilities[keyProbabilities.length - 1];
LogicNode boundsCheck = graph().unique(new IntegerBelowNode(newValue, ConstantNode.forInt(arrayLength, graph())));
graph().addBeforeFixed(this, graph().add(new FixedGuardNode(boundsCheck, DeoptimizationReason.BoundsCheckException, DeoptimizationAction.InvalidateReprofile)));
doReplace(newValue, newKeyDatas, newSuccessors, newDefaultSuccessor, newDefaultProbability);
assert loadIndexed.hasNoUsages();
GraphUtil.removeFixedWithUnusedInputs(loadIndexed);
return true;
}
private static int addNewSuccessor(AbstractBeginNode newSuccessor, ArrayList<AbstractBeginNode> newSuccessors) {
int index = newSuccessors.indexOf(newSuccessor);
if (index == -1) {
index = newSuccessors.size();
newSuccessors.add(newSuccessor);
}
return index;
}
private void doReplace(ValueNode newValue, List<KeyData> newKeyDatas, ArrayList<AbstractBeginNode> newSuccessors, int newDefaultSuccessor, double newDefaultProbability) {
newKeyDatas.sort((k1, k2) -> k1.key - k2.key);
int newKeyCount = newKeyDatas.size();
int[] newKeys = new int[newKeyCount];
double[] newKeyProbabilities = new double[newKeyCount + 1];
int[] newKeySuccessors = new int[newKeyCount + 1];
for (int i = 0; i < newKeyCount; i++) {
KeyData keyData = newKeyDatas.get(i);
newKeys[i] = keyData.key;
newKeyProbabilities[i] = keyData.keyProbability;
newKeySuccessors[i] = keyData.keySuccessor;
}
newKeySuccessors[newKeyCount] = newDefaultSuccessor;
newKeyProbabilities[newKeyCount] = newDefaultProbability;
double totalProbability = 0;
for (double probability : newKeyProbabilities) {
totalProbability += probability;
}
if (totalProbability > 0) {
for (int i = 0; i < newKeyProbabilities.length; i++) {
newKeyProbabilities[i] /= totalProbability;
}
} else {
for (int i = 0; i < newKeyProbabilities.length; i++) {
newKeyProbabilities[i] = 1.0 / newKeyProbabilities.length;
}
}
for (int i = 0; i < blockSuccessorCount(); i++) {
AbstractBeginNode successor = blockSuccessor(i);
if (!newSuccessors.contains(successor)) {
FixedNode fixedBranch = successor;
fixedBranch.predecessor().replaceFirstSuccessor(fixedBranch, null);
GraphUtil.killCFG(fixedBranch);
}
setBlockSuccessor(i, null);
}
AbstractBeginNode[] successorsArray = newSuccessors.toArray(new AbstractBeginNode[newSuccessors.size()]);
SwitchNode newSwitch = graph().add(new IntegerSwitchNode(newValue, successorsArray, newKeys, newKeyProbabilities, newKeySuccessors));
((FixedWithNextNode) predecessor()).setNext(newSwitch);
GraphUtil.killWithUnusedFloatingInputs(this);
}
@Override
public Stamp getValueStampForSuccessor(AbstractBeginNode beginNode) {
Stamp result = null;
if (beginNode != this.defaultSuccessor()) {
for (int i = 0; i < keyCount(); i++) {
if (keySuccessor(i) == beginNode) {
if (result == null) {
result = StampFactory.forConstant(keyAt(i));
} else {
result = result.meet(StampFactory.forConstant(keyAt(i)));
}
}
}
}
return result;
}
}