package codechicken.lib.asm;
import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;
import com.google.common.collect.ImmutableMap;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.JumpInsnNode;
import org.objectweb.asm.tree.LabelNode;
import java.util.*;
import java.util.Map.Entry;
import static org.objectweb.asm.tree.AbstractInsnNode.*;
public class ASMBlock
{
public InsnListSection list;
private BiMap<String, LabelNode> labels;
public ASMBlock(InsnListSection list, BiMap<String, LabelNode> labels) {
this.list = list;
this.labels = labels;
}
public ASMBlock(InsnListSection list) {
this(list, HashBiMap.<String, LabelNode>create());
}
public ASMBlock(InsnList list) {
this(new InsnListSection(list));
}
public ASMBlock() {
this(new InsnListSection());
}
public LabelNode getOrAdd(String s) {
LabelNode l = get(s);
if (l == null)
labels.put(s, l = new LabelNode());
return l;
}
public LabelNode get(String s) {
return labels.get(s);
}
public void replaceLabels(Map<LabelNode, LabelNode> labelMap, Set<LabelNode> usedLabels) {
for (AbstractInsnNode insn : list)
switch (insn.getType()) {
case LABEL:
AbstractInsnNode insn2 = insn.clone(labelMap);
if (insn2 == insn)//identity mapping
continue;
if (usedLabels.contains(insn2))
throw new IllegalStateException("LabelNode cannot be a part of two InsnLists");
list.replace(insn, insn2);
break;
case JUMP_INSN:
case FRAME:
case LOOKUPSWITCH_INSN:
case TABLESWITCH_INSN:
list.replace(insn, insn.clone(labelMap));
}
for(Entry<LabelNode, LabelNode> entry : labelMap.entrySet()) {
String key = labels.inverse().get(entry.getKey());
if(key != null)
labels.put(key, entry.getValue());
}
}
public void replaceLabels(Map<LabelNode, LabelNode> labelMap) {
replaceLabels(labelMap, Collections.EMPTY_SET);
}
public void replaceLabel(String s, LabelNode l) {
LabelNode old = get(s);
if (old != null)
replaceLabels(ImmutableMap.of(old, l));
}
/**
* Pulls all common labels from other into this
* @return this
*/
public ASMBlock mergeLabels(ASMBlock other) {
if(labels.isEmpty() || other.labels.isEmpty())
return this;
//common labels, give them our nodes
HashMap<LabelNode, LabelNode> labelMap = list.identityLabelMap();
for(Entry<String, LabelNode> entry : other.labels.entrySet()) {
LabelNode old = labels.get(entry.getKey());
if(old != null)
labelMap.put(old, entry.getValue());
}
HashSet<LabelNode> usedLabels = new HashSet<LabelNode>();
for (AbstractInsnNode insn = other.list.list.getFirst(); insn != null; insn = insn.getNext())
if(insn.getType() == LABEL)
usedLabels.add((LabelNode) insn);
replaceLabels(labelMap, usedLabels);
return this;
}
/**
* Like mergeLabels but pulls insns from other list into this so LabelNodes can be transferred
* @return this
*/
public ASMBlock pullLabels(ASMBlock other) {
other.list.remove();
return mergeLabels(other);
}
public ASMBlock copy() {
BiMap<String, LabelNode> labels = HashBiMap.create();
Map<LabelNode, LabelNode> labelMap = list.cloneLabels();
for(Entry<String, LabelNode> entry : this.labels.entrySet())
labels.put(entry.getKey(), labelMap.get(entry.getValue()));
return new ASMBlock(list.copy(labelMap), labels);
}
public ASMBlock applyLabels(InsnListSection list2) {
if(labels.isEmpty())
return new ASMBlock(list2);
Set<LabelNode> cFlowLabels1 = labels.values();
Set<LabelNode> cFlowLabels2 = InsnComparator.getControlFlowLabels(list2);
ASMBlock block = new ASMBlock(list2);
HashMap<LabelNode, LabelNode> labelMap = new HashMap<LabelNode, LabelNode>();
for(int i = 0, k = 0; i < list.size() && k < list2.size(); ) {
AbstractInsnNode insn1 = list.get(i);
if(!InsnComparator.insnImportant(insn1, cFlowLabels1)) {
i++;
continue;
}
AbstractInsnNode insn2 = list2.get(k);
if(!InsnComparator.insnImportant(insn2, cFlowLabels2)) {
k++;
continue;
}
if(insn1.getOpcode() != insn2.getOpcode())
throw new IllegalArgumentException("Lists do not match:\n"+list+"\n\n"+list2);
switch(insn1.getType()) {
case LABEL:
labelMap.put((LabelNode) insn1, (LabelNode) insn2);
break;
case JUMP_INSN:
labelMap.put(((JumpInsnNode) insn1).label, ((JumpInsnNode) insn2).label);
break;
}
i++; k++;
}
for(Entry<String, LabelNode> entry : labels.entrySet())
block.labels.put(entry.getKey(), labelMap.get(entry.getValue()));
return block;
}
public InsnList rawListCopy() {
return list.copy().list;
}
}