Commit 46c85728 authored by Skylot's avatar Skylot

core: restore switch over enum

parent e6b91900
......@@ -10,6 +10,7 @@ import jadx.core.dex.instructions.SwitchNode;
import jadx.core.dex.instructions.args.InsnArg;
import jadx.core.dex.instructions.args.RegisterArg;
import jadx.core.dex.nodes.BlockNode;
import jadx.core.dex.nodes.FieldNode;
import jadx.core.dex.nodes.IBlock;
import jadx.core.dex.nodes.IContainer;
import jadx.core.dex.nodes.IRegion;
......@@ -239,7 +240,14 @@ public class RegionGen extends InsnGen {
IContainer c = sw.getCases().get(i);
for (Object k : keys) {
code.startLine("case ");
if (k instanceof IndexInsnNode) {
if (k instanceof FieldNode) {
FieldNode fn = (FieldNode) k;
if (fn.getParentClass().isEnum()) {
code.add(fn.getName());
} else {
staticField(code, fn.getFieldInfo());
}
} else if (k instanceof IndexInsnNode) {
staticField(code, (FieldInfo) ((IndexInsnNode) k).getIndex());
} else {
code.add(TypeGen.literalToString((Integer) k, arg.getType()));
......
......@@ -4,6 +4,7 @@ import jadx.core.dex.attributes.annotations.AnnotationsList;
import jadx.core.dex.attributes.annotations.MethodParameters;
import jadx.core.dex.attributes.nodes.DeclareVariablesAttr;
import jadx.core.dex.attributes.nodes.EnumClassAttr;
import jadx.core.dex.attributes.nodes.EnumMapAttr;
import jadx.core.dex.attributes.nodes.FieldReplaceAttr;
import jadx.core.dex.attributes.nodes.ForceReturnAttr;
import jadx.core.dex.attributes.nodes.JadxErrorAttr;
......@@ -40,6 +41,7 @@ public class AType<T extends IAttribute> {
public static final AType<JadxErrorAttr> JADX_ERROR = new AType<JadxErrorAttr>();
public static final AType<MethodInlineAttr> METHOD_INLINE = new AType<MethodInlineAttr>();
public static final AType<EnumClassAttr> ENUM_CLASS = new AType<EnumClassAttr>();
public static final AType<EnumMapAttr> ENUM_MAP = new AType<EnumMapAttr>();
public static final AType<AnnotationsList> ANNOTATION_LIST = new AType<AnnotationsList>();
public static final AType<MethodParameters> ANNOTATION_MTH_PARAMETERS = new AType<MethodParameters>();
public static final AType<PhiListAttr> PHI_LIST = new AType<PhiListAttr>();
......
......@@ -12,25 +12,25 @@ public abstract class AttrNode implements IAttributeNode {
@Override
public void add(AFlag flag) {
getStorage().add(flag);
initStorage().add(flag);
}
@Override
public void addAttr(IAttribute attr) {
getStorage().add(attr);
initStorage().add(attr);
}
@Override
public <T> void addAttr(AType<AttrList<T>> type, T obj) {
getStorage().add(type, obj);
initStorage().add(type, obj);
}
@Override
public void copyAttributesFrom(AttrNode attrNode) {
getStorage().addAll(attrNode.storage);
initStorage().addAll(attrNode.storage);
}
AttributeStorage getStorage() {
AttributeStorage initStorage() {
AttributeStorage store = storage;
if (store == EMPTY_ATTR_STORAGE) {
store = new AttributeStorage();
......
......@@ -5,7 +5,7 @@ import jadx.core.dex.attributes.annotations.Annotation;
import java.util.Collections;
import java.util.List;
public class EmptyAttrStorage extends AttributeStorage {
public final class EmptyAttrStorage extends AttributeStorage {
@Override
public boolean contains(AFlag flag) {
......@@ -52,4 +52,9 @@ public class EmptyAttrStorage extends AttributeStorage {
public List<String> getAttributeStrings() {
return Collections.emptyList();
}
@Override
public String toString() {
return "";
}
}
package jadx.core.dex.attributes.nodes;
import jadx.core.dex.attributes.AType;
import jadx.core.dex.attributes.IAttribute;
import java.util.HashMap;
import java.util.Map;
public class EnumMapAttr implements IAttribute {
private Map<Object, Object> map = new HashMap<Object, Object>();
public Map<Object, Object> getMap() {
return map;
}
@Override
public AType<EnumMapAttr> getType() {
return AType.ENUM_MAP;
}
@Override
public String toString() {
return "Enum fields map: " + map;
}
}
......@@ -38,13 +38,12 @@ public class EnumVisitor extends AbstractVisitor {
// collect enum fields, remove synthetic
List<FieldNode> enumFields = new ArrayList<FieldNode>();
for (Iterator<FieldNode> it = cls.getFields().iterator(); it.hasNext(); ) {
FieldNode f = it.next();
for (FieldNode f : cls.getFields()) {
if (f.getAccessFlags().isEnum()) {
enumFields.add(f);
it.remove();
f.add(AFlag.DONT_GENERATE);
} else if (f.getAccessFlags().isSynthetic()) {
it.remove();
f.add(AFlag.DONT_GENERATE);
}
}
......
package jadx.core.dex.visitors;
import jadx.core.dex.attributes.AFlag;
import jadx.core.dex.attributes.AType;
import jadx.core.dex.attributes.nodes.EnumMapAttr;
import jadx.core.dex.info.FieldInfo;
import jadx.core.dex.instructions.IndexInsnNode;
import jadx.core.dex.instructions.InsnType;
import jadx.core.dex.instructions.InvokeNode;
import jadx.core.dex.instructions.SwitchNode;
import jadx.core.dex.instructions.args.InsnArg;
import jadx.core.dex.instructions.args.InsnWrapArg;
import jadx.core.dex.instructions.args.LiteralArg;
import jadx.core.dex.nodes.BlockNode;
import jadx.core.dex.nodes.ClassNode;
import jadx.core.dex.nodes.FieldNode;
import jadx.core.dex.nodes.InsnNode;
import jadx.core.dex.nodes.MethodNode;
import jadx.core.utils.InstructionRemover;
import jadx.core.utils.exceptions.DecodeException;
import jadx.core.utils.exceptions.JadxException;
import java.util.List;
......@@ -43,6 +54,9 @@ public class ReSugarCode extends AbstractVisitor {
case NEW_ARRAY:
return processNewArray(mth, instructions, i, remover);
case SWITCH:
return processEnumSwitch(mth, (SwitchNode) insn);
default:
return null;
}
......@@ -75,4 +89,153 @@ public class ReSugarCode extends AbstractVisitor {
}
return filledArr;
}
private static InsnNode processEnumSwitch(MethodNode mth, SwitchNode insn) {
InsnArg arg = insn.getArg(0);
if (!arg.isInsnWrap()) {
return null;
}
InsnNode wrapInsn = ((InsnWrapArg) arg).getWrapInsn();
if (wrapInsn.getType() != InsnType.AGET) {
return null;
}
EnumMapInfo enumMapInfo = checkEnumMapAccess(mth, wrapInsn);
if (enumMapInfo == null) {
return null;
}
FieldNode enumMapField = enumMapInfo.getMapField();
InsnArg invArg = enumMapInfo.getArg();
EnumMapAttr enumMapAttr = getEnumMap(mth, enumMapField);
if (enumMapAttr == null) {
return null;
}
Object[] keys = insn.getKeys();
for (int i = 0; i < keys.length; i++) {
Object key = keys[i];
Object newKey = enumMapAttr.getMap().get(key);
if (newKey != null) {
keys[i] = newKey;
} else {
return null;
}
}
enumMapField.getParentClass().add(AFlag.DONT_GENERATE);
insn.replaceArg(arg, invArg);
return null;
}
private static EnumMapAttr getEnumMap(MethodNode mth, FieldNode field) {
ClassNode syntheticClass = field.getParentClass();
EnumMapAttr mapAttr = syntheticClass.get(AType.ENUM_MAP);
if (mapAttr != null) {
return mapAttr;
}
MethodNode clsInitMth = syntheticClass.searchMethodByName("<clinit>()V");
if (clsInitMth == null || clsInitMth.isNoCode()) {
return null;
}
if (clsInitMth.getBasicBlocks() == null) {
try {
clsInitMth.load();
} catch (DecodeException e) {
LOG.error("Load failed", e);
return null;
}
if (clsInitMth.getBasicBlocks() == null) {
// TODO:
return null;
}
}
mapAttr = new EnumMapAttr();
for (BlockNode block : clsInitMth.getBasicBlocks()) {
for (InsnNode insn : block.getInstructions()) {
if (insn.getType() == InsnType.APUT) {
addToEnumMap(mth, field, mapAttr, insn);
}
}
}
if (mapAttr.getMap().isEmpty()) {
return null;
}
syntheticClass.addAttr(mapAttr);
return mapAttr;
}
private static void addToEnumMap(MethodNode mth, FieldNode field, EnumMapAttr mapAttr, InsnNode insn) {
InsnArg litArg = insn.getArg(2);
if (!litArg.isLiteral()) {
return;
}
EnumMapInfo mapInfo = checkEnumMapAccess(mth, insn);
if (mapInfo == null) {
return;
}
InsnArg enumArg = mapInfo.getArg();
if (mapInfo.getMapField() != field || !enumArg.isInsnWrap()) {
return;
}
InsnNode sget = ((InsnWrapArg) enumArg).getWrapInsn();
if (!(sget instanceof IndexInsnNode)) {
return;
}
Object index = ((IndexInsnNode) sget).getIndex();
if (!(index instanceof FieldInfo)) {
return;
}
FieldNode fieldNode = mth.dex().resolveField((FieldInfo) index);
if (fieldNode == null) {
return;
}
int literal = (int) ((LiteralArg) litArg).getLiteral();
mapAttr.getMap().put(literal, fieldNode);
}
public static EnumMapInfo checkEnumMapAccess(MethodNode mth, InsnNode checkInsn) {
InsnArg sgetArg = checkInsn.getArg(0);
InsnArg invArg = checkInsn.getArg(1);
if (!sgetArg.isInsnWrap() || !invArg.isInsnWrap()) {
return null;
}
InsnNode invInsn = ((InsnWrapArg) invArg).getWrapInsn();
InsnNode sgetInsn = ((InsnWrapArg) sgetArg).getWrapInsn();
if (invInsn.getType() != InsnType.INVOKE || sgetInsn.getType() != InsnType.SGET) {
return null;
}
InvokeNode inv = (InvokeNode) invInsn;
if (!inv.getCallMth().getShortId().equals("ordinal()I")) {
return null;
}
ClassNode enumCls = mth.dex().resolveClass(inv.getCallMth().getDeclClass());
if (enumCls == null || !enumCls.isEnum()) {
return null;
}
Object index = ((IndexInsnNode) sgetInsn).getIndex();
if (!(index instanceof FieldInfo)) {
return null;
}
FieldNode enumMapField = mth.dex().resolveField(((FieldInfo) index));
if (enumMapField == null || !enumMapField.getAccessFlags().isSynthetic()) {
return null;
}
return new EnumMapInfo(inv.getArg(0), enumMapField);
}
private static class EnumMapInfo {
private final InsnArg arg;
private final FieldNode mapField;
public EnumMapInfo(InsnArg arg, FieldNode mapField) {
this.arg = arg;
this.mapField = mapField;
}
public InsnArg getArg() {
return arg;
}
public FieldNode getMapField() {
return mapField;
}
}
}
package jadx.tests.internal.enums;
import jadx.api.InternalJadxTest;
import jadx.core.dex.nodes.ClassNode;
import org.junit.Test;
import static jadx.tests.utils.JadxMatchers.countString;
import static org.junit.Assert.assertThat;
public class TestSwitchOverEnum extends InternalJadxTest {
public enum Count {
ONE, TWO, THREE
}
public int testEnum(Count c) {
switch (c) {
case ONE:
return 1;
case TWO:
return 2;
case THREE:
return 3;
}
return 0;
}
@Test
public void test() {
ClassNode cls = getClassNode(TestSwitchOverEnum.class);
String code = cls.getCode().toString();
System.out.println(code);
assertThat(code, countString(1, "synthetic"));
assertThat(code, countString(2, "switch (c) {"));
assertThat(code, countString(2, "case ONE:"));
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment