core: fix switch over enum with several enums in class

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

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

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

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

jadx-core/src/test/java/jadx/tests/integration/enums/TestSwitchOverEnum2.java

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