package me.nov.cafebabe.analysis.blocks;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.JumpInsnNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.LookupSwitchInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TableSwitchInsnNode;
import me.nov.cafebabe.utils.formatting.OpcodeFormatting;
public class Converter implements Opcodes {
private ArrayList<AbstractInsnNode> nodes;
private static final boolean DEBUG = false;
public Converter(MethodNode mn) {
assert (mn.instructions != null && mn.instructions.size() > 0);
this.nodes = new ArrayList<>(Arrays.asList(mn.instructions.toArray()));
}
public Converter(AbstractInsnNode[] array) {
assert (array != null && array.length > 0);
this.nodes = new ArrayList<>(Arrays.asList(array));
}
public ArrayList<Block> convert(boolean simplify, boolean removeRedundant, boolean skipDupeSwitches, int maxInputRemoveNonsense) {
ArrayList<Block> blocks = new ArrayList<>();
HashMap<AbstractInsnNode, Block> correspBlock = new HashMap<>();
Block block = null;
if (nodes.isEmpty()) {
return blocks;
}
//detect block structure & last block insns
for (AbstractInsnNode ain : nodes) {
if (block == null) {
block = new Block();
}
if (ain instanceof LabelNode) {
block.setLabel((LabelNode) ain);
}
block.getNodes().add(ain);
correspBlock.put(ain, block);
//end blocks
int op = ain.getOpcode();
if (op >= IRETURN && op <= RETURN || ain instanceof JumpInsnNode || op == ATHROW || op == LOOKUPSWITCH || op == TABLESWITCH) {
block.setEndNode(ain);
blocks.add(block);
block = null;
continue;
}
//next because the label should have a new block
//blocks that end without a jump
if (ain.getNext() != null && (ain.getNext() instanceof LabelNode)) {
block.setEndNode(ain.getNext());
blocks.add(block);
block = null;
continue;
}
}
for (Block b : blocks) {
AbstractInsnNode end = b.getEndNode();
//only opc where it continues
if (end instanceof JumpInsnNode) {
JumpInsnNode jin = (JumpInsnNode) end;
ArrayList<Block> outputs = new ArrayList<>();
if (!correspBlock.containsKey(jin.label)) {
throw new RuntimeException("label not visited");
}
Block blockAtLabel = correspBlock.get(jin.label);
//blockAtLabel can be the same block!
if (end.getOpcode() == GOTO) {
outputs.add(blockAtLabel);
b.setOutput(outputs);
blockAtLabel.getInput().add(b);
} else {
//ifs have two outputs: either it jumps or not
outputs.add(blockAtLabel);
if (jin.getNext() == null) {
throw new RuntimeException("if has no next entry");
}
if (correspBlock.get(jin.getNext()) == b) {
throw new RuntimeException("next node is self?");
}
Block blockAfter = correspBlock.get(jin.getNext());
outputs.add(blockAfter);
b.setOutput(outputs);
blockAtLabel.getInput().add(b);
blockAfter.getInput().add(b);
}
} else if (end instanceof TableSwitchInsnNode) {
ArrayList<Block> outputs = new ArrayList<>();
TableSwitchInsnNode tsin = (TableSwitchInsnNode) end;
if (tsin.dflt != null) {
Block blockAtDefault = correspBlock.get(tsin.dflt);
blockAtDefault.getInput().add(b);
outputs.add(blockAtDefault);
}
ArrayList<LabelNode> alreadyConnected = new ArrayList<>();
for (LabelNode l : tsin.labels) {
if (skipDupeSwitches) {
if (alreadyConnected.contains(l))
continue;
alreadyConnected.add(l);
}
Block blockAtCase = correspBlock.get(l);
blockAtCase.getInput().add(b);
outputs.add(blockAtCase);
}
b.setOutput(outputs);
} else if (end instanceof LookupSwitchInsnNode) {
ArrayList<Block> outputs = new ArrayList<>();
LookupSwitchInsnNode lsin = (LookupSwitchInsnNode) end;
if (lsin.dflt != null) {
Block blockAtDefault = correspBlock.get(lsin.dflt);
blockAtDefault.getInput().add(b);
outputs.add(blockAtDefault);
}
ArrayList<LabelNode> alreadyConnected = new ArrayList<>();
for (LabelNode l : lsin.labels) {
if (skipDupeSwitches) {
if (alreadyConnected.contains(l))
continue;
alreadyConnected.add(l);
}
Block blockAtCase = correspBlock.get(l);
blockAtCase.getInput().add(b);
outputs.add(blockAtCase);
}
b.setOutput(outputs);
} else if (end instanceof LabelNode) {
if (!correspBlock.containsKey(end)) {
throw new RuntimeException("label not visited");
}
ArrayList<Block> outputs = new ArrayList<>();
Block blockAtNext = correspBlock.get(end);
outputs.add(blockAtNext);
b.setOutput(outputs);
blockAtNext.getInput().add(b);
}
}
Block first = correspBlock.get(nodes.get(0));
assert (first != null);
if (removeRedundant) {
ArrayList<Block> visited = new ArrayList<>();
removeNonsense(visited, blocks, first, maxInputRemoveNonsense);
for (Block b : new ArrayList<>(blocks)) {
if (b.getInput().isEmpty()) {
removeNonsense(visited, blocks, b, maxInputRemoveNonsense);
}
}
}
if (simplify) {
ArrayList<Block> visited = new ArrayList<>();
simplifyBlock(visited, blocks, first);
for (Block b : new ArrayList<>(blocks)) {
if (b.getInput().isEmpty()) {
simplifyBlock(visited, blocks, b);
}
}
}
if (DEBUG) {
ArrayList<Block> visited = new ArrayList<>();
calculateDepths(visited, blocks, first, 0);
for (Block b : blocks) {
if (b.getInput().isEmpty()) {
calculateDepths(visited, blocks, b, 0);
}
}
}
return blocks;
}
private void calculateDepths(ArrayList<Block> visited, ArrayList<Block> blocks, Block b, int depth) {
if (visited.contains(b)) {
return;
}
visited.add(b);
b.setDepth(depth);
if (b.endsWithJump()) {
if (b.getOutput().size() > 1) {
if (b.getInput().size() <= 1) {
//handle if surrounding
Block blockAtLabel = b.getOutput().get(0);
Block blockAfter = b.getOutput().get(1);
calculateDepths(visited, blocks, blockAtLabel, depth);
calculateDepths(visited, blocks, blockAfter, depth + 1);
} else {
//handle while surrounding
Block blockAtLabel = b.getOutput().get(0);
Block blockAfter = b.getOutput().get(1);
calculateDepths(visited, blocks, blockAfter, depth);
calculateDepths(visited, blocks, blockAtLabel, depth + 1);
}
}
}
for (Block output : b.getOutput()) {
calculateDepths(visited, blocks, output, depth);
}
}
private void removeNonsense(ArrayList<Block> visited, ArrayList<Block> blocks, Block b, int maxInputRemoveNonsense) {
if (visited.contains(b)) {
return;
}
visited.add(b);
if (b.endsWithJump()) {
if (b.getInput().size() <= maxInputRemoveNonsense && b.getOutput().size() == 1) { //there could be more inputs but that might lead to a unreadable graph
if (isJumpBlock(b)) {
Block output = b.getOutput().get(0);
for (Block input : b.getInput()) {
input.getOutput().remove(b);
input.getOutput().add(output);
output.getInput().add(input);
}
output.getInput().remove(b);
}
}
}
for (Block output : new ArrayList<>(b.getOutput())) {
removeNonsense(visited, blocks, output, maxInputRemoveNonsense);
}
}
private boolean isJumpBlock(Block b) {
for (AbstractInsnNode ain : b.getNodes()) {
int type = ain.getType();
if (type != AbstractInsnNode.LABEL && type != AbstractInsnNode.LINE && type != AbstractInsnNode.FRAME && type != AbstractInsnNode.JUMP_INSN) {
return false;
}
}
return true;
}
private void simplifyBlock(ArrayList<Block> simplified, ArrayList<Block> blocks, Block b) {
if (simplified.contains(b)) {
return;
}
simplified.add(b);
while (true) {
if (b.getOutput().size() == 1) {
Block to = b.getOutput().get(0);
//also optimizes unnecessary gotos
if (to.getInput().size() == 1 && !isFirst(to)) {
assert (to.getInput().get(0) == b);
b.getNodes().addAll(to.getNodes());
b.setEndNode(to.getEndNode());
b.setOutput(to.getOutput());
blocks.remove(to);
continue;
}
}
break;
}
for (Block output : b.getOutput()) {
simplifyBlock(simplified, blocks, output);
}
}
private boolean isFirst(Block to) {
LabelNode ln = to.getLabel();
if (ln != null && OpcodeFormatting.getLabelIndex(ln) == 0) {
return true;
}
return false;
}
}
