core: restore switch over enum

jadx-core/src/main/java/jadx/core/codegen/RegionGen.java

@@ -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()));

jadx-core/src/main/java/jadx/core/dex/attributes/AType.java

@@ -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>();

jadx-core/src/main/java/jadx/core/dex/attributes/AttrNode.java

@@ -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();

jadx-core/src/main/java/jadx/core/dex/attributes/EmptyAttrStorage.java

@@ -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 "";
+	}
 }

jadx-core/src/main/java/jadx/core/dex/attributes/nodes/EnumMapAttr.java

+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;
+	}
+
+}

jadx-core/src/main/java/jadx/core/dex/visitors/EnumVisitor.java

@@ -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);
 			}
 		}
 

jadx-core/src/main/java/jadx/core/dex/visitors/ReSugarCode.java

 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;
+		}
+	}
 }

jadx-core/src/test/java/jadx/tests/internal/enums/TestSwitchOverEnum.java

+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:"));
+	}
+}