Commit 84970759 authored by Skylot's avatar Skylot

core: fix switch over enum with several enums in class

parent 53cac58e
......@@ -2,16 +2,38 @@ package jadx.core.dex.attributes.nodes;
import jadx.core.dex.attributes.AType;
import jadx.core.dex.attributes.IAttribute;
import jadx.core.dex.nodes.FieldNode;
import java.util.HashMap;
import java.util.Map;
public class EnumMapAttr implements IAttribute {
private Map<Object, Object> map = new HashMap<Object, Object>();
public static class KeyValueMap {
private Map<Object, Object> map = new HashMap<Object, Object>();
public Map<Object, Object> getMap() {
return map;
public Object get(Object key) {
return map.get(key);
}
void put(Object key, Object value) {
map.put(key, value);
}
}
private Map<FieldNode, KeyValueMap> fieldsMap = new HashMap<FieldNode, KeyValueMap>();
public KeyValueMap getMap(FieldNode field) {
return fieldsMap.get(field);
}
public void add(FieldNode field, Object key, Object value) {
KeyValueMap map = getMap(field);
if (map == null) {
map = new KeyValueMap();
fieldsMap.put(field, map);
}
map.put(key, value);
}
@Override
......@@ -21,7 +43,7 @@ public class EnumMapAttr implements IAttribute {
@Override
public String toString() {
return "Enum fields map: " + map;
return "Enum fields map: " + fieldsMap;
}
}
......@@ -3,6 +3,7 @@ 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.AccessInfo;
import jadx.core.dex.info.FieldInfo;
import jadx.core.dex.instructions.IndexInsnNode;
import jadx.core.dex.instructions.InsnType;
......@@ -106,31 +107,38 @@ public class ReSugarCode extends AbstractVisitor {
FieldNode enumMapField = enumMapInfo.getMapField();
InsnArg invArg = enumMapInfo.getArg();
EnumMapAttr enumMapAttr = getEnumMap(mth, enumMapField);
if (enumMapAttr == null) {
EnumMapAttr.KeyValueMap valueMap = getEnumMap(mth, enumMapField);
if (valueMap == 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 {
for (Object key : keys) {
Object newKey = valueMap.get(key);
if (newKey == null) {
return null;
}
}
enumMapField.getParentClass().add(AFlag.DONT_GENERATE);
insn.replaceArg(arg, invArg);
// replace confirmed
if (!insn.replaceArg(arg, invArg)) {
return null;
}
for (int i = 0; i < keys.length; i++) {
keys[i] = valueMap.get(keys[i]);
}
enumMapField.add(AFlag.DONT_GENERATE);
checkAndHideClass(enumMapField.getParentClass());
return null;
}
private static EnumMapAttr getEnumMap(MethodNode mth, FieldNode field) {
private static EnumMapAttr.KeyValueMap getEnumMap(MethodNode mth, FieldNode field) {
ClassNode syntheticClass = field.getParentClass();
EnumMapAttr mapAttr = syntheticClass.get(AType.ENUM_MAP);
if (mapAttr != null) {
return mapAttr;
return mapAttr.getMap(field);
}
mapAttr = new EnumMapAttr();
syntheticClass.addAttr(mapAttr);
MethodNode clsInitMth = syntheticClass.searchMethodByName("<clinit>()V");
if (clsInitMth == null || clsInitMth.isNoCode()) {
return null;
......@@ -147,32 +155,28 @@ public class ReSugarCode extends AbstractVisitor {
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);
addToEnumMap(mth, mapAttr, insn);
}
}
}
if (mapAttr.getMap().isEmpty()) {
return null;
}
syntheticClass.addAttr(mapAttr);
return mapAttr;
return mapAttr.getMap(field);
}
private static void addToEnumMap(MethodNode mth, FieldNode field, EnumMapAttr mapAttr, InsnNode insn) {
InsnArg litArg = insn.getArg(2);
private static void addToEnumMap(MethodNode mth, EnumMapAttr mapAttr, InsnNode aputInsn) {
InsnArg litArg = aputInsn.getArg(2);
if (!litArg.isLiteral()) {
return;
}
EnumMapInfo mapInfo = checkEnumMapAccess(mth, insn);
EnumMapInfo mapInfo = checkEnumMapAccess(mth, aputInsn);
if (mapInfo == null) {
return;
}
InsnArg enumArg = mapInfo.getArg();
if (mapInfo.getMapField() != field || !enumArg.isInsnWrap()) {
FieldNode field = mapInfo.getMapField();
if (field == null || !enumArg.isInsnWrap()) {
return;
}
InsnNode sget = ((InsnWrapArg) enumArg).getWrapInsn();
......@@ -188,7 +192,7 @@ public class ReSugarCode extends AbstractVisitor {
return;
}
int literal = (int) ((LiteralArg) litArg).getLiteral();
mapAttr.getMap().put(literal, fieldNode);
mapAttr.add(field, literal, fieldNode);
}
public static EnumMapInfo checkEnumMapAccess(MethodNode mth, InsnNode checkInsn) {
......@@ -221,6 +225,20 @@ public class ReSugarCode extends AbstractVisitor {
return new EnumMapInfo(inv.getArg(0), enumMapField);
}
/**
* If all static final synthetic fields have DONT_GENERATE => hide whole class
*/
private static void checkAndHideClass(ClassNode cls) {
for (FieldNode field : cls.getFields()) {
AccessInfo af = field.getAccessFlags();
if (af.isSynthetic() && af.isStatic() && af.isFinal()
&& !field.contains(AFlag.DONT_GENERATE)) {
return;
}
}
cls.add(AFlag.DONT_GENERATE);
}
private static class EnumMapInfo {
private final InsnArg arg;
private final FieldNode mapField;
......
package jadx.tests.integration.enums;
import jadx.core.dex.nodes.ClassNode;
import jadx.tests.api.IntegrationTest;
import org.junit.Test;
import static jadx.tests.api.utils.JadxMatchers.countString;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
public class TestSwitchOverEnum2 extends IntegrationTest {
public enum Count {
ONE, TWO, THREE
}
public enum Animal {
CAT, DOG
}
public int testEnum(Count c, Animal a) {
int result = 0;
switch (c) {
case ONE:
result = 1;
break;
case TWO:
result = 2;
break;
}
switch (a) {
case CAT:
result += 10;
break;
case DOG:
result += 20;
break;
}
return result;
}
public void check() {
assertEquals(21, testEnum(Count.ONE, Animal.DOG));
}
@Test
public void test() {
ClassNode cls = getClassNode(TestSwitchOverEnum2.class);
String code = cls.getCode().toString();
assertThat(code, countString(1, "synthetic"));
assertThat(code, countString(2, "switch (c) {"));
assertThat(code, countString(2, "case ONE:"));
assertThat(code, countString(2, "case DOG:"));
}
}
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