core: move enum restore pass to later stage

jadx-core/src/main/java/jadx/core/Jadx.java

@@ -79,7 +79,6 @@ public class Jadx {
 			passes.add(new EliminatePhiNodes());
 
 			passes.add(new ModVisitor());
-			passes.add(new EnumVisitor());
 
 			passes.add(new CodeShrinker());
 			passes.add(new ReSugarCode());
@@ -102,6 +101,7 @@ public class Jadx {
 
 			passes.add(new MethodInlineVisitor());
 			passes.add(new ClassModifier());
+			passes.add(new EnumVisitor());
 			passes.add(new PrepareForCodeGen());
 			passes.add(new LoopRegionVisitor());
 			passes.add(new ProcessVariables());

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

@@ -10,8 +10,8 @@ import jadx.core.dex.attributes.nodes.SourceFileAttr;
 import jadx.core.dex.info.AccessInfo;
 import jadx.core.dex.info.ClassInfo;
 import jadx.core.dex.instructions.args.ArgType;
-import jadx.core.dex.instructions.args.InsnArg;
 import jadx.core.dex.instructions.args.PrimitiveType;
+import jadx.core.dex.instructions.mods.ConstructorInsn;
 import jadx.core.dex.nodes.ClassNode;
 import jadx.core.dex.nodes.DexNode;
 import jadx.core.dex.nodes.FieldNode;
@@ -371,21 +371,14 @@ public class ClassGen {
 		for (Iterator<EnumField> it = enumFields.getFields().iterator(); it.hasNext(); ) {
 			EnumField f = it.next();
 			code.startLine(f.getName());
-			if (!f.getArgs().isEmpty()) {
-				code.add('(');
-				for (Iterator<InsnArg> aIt = f.getArgs().iterator(); aIt.hasNext(); ) {
-					InsnArg arg = aIt.next();
-					if (igen == null) {
-						// don't init mth gen if this is simple enum
-						MethodGen mthGen = new MethodGen(this, enumFields.getStaticMethod());
-						igen = new InsnGen(mthGen, false);
-					}
-					igen.addArg(code, arg);
-					if (aIt.hasNext()) {
-						code.add(", ");
-					}
+			ConstructorInsn constrInsn = f.getConstrInsn();
+			if (constrInsn.getArgsCount() > f.getStartArg()) {
+				if (igen == null) {
+					MethodGen mthGen = new MethodGen(this, enumFields.getStaticMethod());
+					igen = new InsnGen(mthGen, false);
 				}
-				code.add(')');
+				MethodNode callMth = cls.dex().resolveMethod(constrInsn.getCallMth());
+				igen.generateMethodArguments(code, constrInsn, f.getStartArg(), callMth);
 			}
 			if (f.getCls() != null) {
 				code.add(' ');

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

@@ -616,7 +616,7 @@ public class InsnGen {
 		generateMethodArguments(code, insn, k, callMthNode);
 	}
 
-	private void generateMethodArguments(CodeWriter code, InsnNode insn, int startArgNum,
+	void generateMethodArguments(CodeWriter code, InsnNode insn, int startArgNum,
 			@Nullable MethodNode callMth) throws CodegenException {
 		int k = startArgNum;
 		if (callMth != null && callMth.contains(AFlag.SKIP_FIRST_ARG)) {

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

@@ -2,37 +2,37 @@ package jadx.core.dex.attributes.nodes;
 
 import jadx.core.dex.attributes.AType;
 import jadx.core.dex.attributes.IAttribute;
-import jadx.core.dex.instructions.args.InsnArg;
+import jadx.core.dex.instructions.mods.ConstructorInsn;
 import jadx.core.dex.nodes.ClassNode;
 import jadx.core.dex.nodes.MethodNode;
-import jadx.core.utils.Utils;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
 public class EnumClassAttr implements IAttribute {
 
 	public static class EnumField {
 		private final String name;
-		private final List<InsnArg> args;
+		private final ConstructorInsn constrInsn;
+		private final int startArg;
 		private ClassNode cls;
 
-		public EnumField(String name, int argsCount) {
+		public EnumField(String name, ConstructorInsn co, int startArg) {
 			this.name = name;
-			if (argsCount != 0) {
-				this.args = new ArrayList<InsnArg>(argsCount);
-			} else {
-				this.args = Collections.emptyList();
-			}
+			this.constrInsn = co;
+			this.startArg = startArg;
 		}
 
 		public String getName() {
 			return name;
 		}
 
-		public List<InsnArg> getArgs() {
-			return args;
+		public ConstructorInsn getConstrInsn() {
+			return constrInsn;
+		}
+
+		public int getStartArg() {
+			return startArg;
 		}
 
 		public ClassNode getCls() {
@@ -45,7 +45,7 @@ public class EnumClassAttr implements IAttribute {
 
 		@Override
 		public String toString() {
-			return name + "(" + Utils.listToString(args) + ") " + cls;
+			return name + "(" + constrInsn + ") " + cls;
 		}
 	}
 

jadx-core/src/main/java/jadx/core/dex/instructions/args/RegisterArg.java

 package jadx.core.dex.instructions.args;
 
-import jadx.core.dex.attributes.AType;
-import jadx.core.dex.info.FieldInfo;
-import jadx.core.dex.instructions.ConstClassNode;
-import jadx.core.dex.instructions.ConstStringNode;
-import jadx.core.dex.instructions.IndexInsnNode;
 import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.instructions.PhiInsn;
 import jadx.core.dex.nodes.DexNode;
-import jadx.core.dex.nodes.FieldNode;
 import jadx.core.dex.nodes.InsnNode;
-import jadx.core.dex.nodes.parser.FieldValueAttr;
+import jadx.core.utils.InsnUtils;
 
 import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
@@ -98,28 +92,7 @@ public class RegisterArg extends InsnArg implements Named {
 		if (parInsn == null) {
 			return null;
 		}
-		InsnType insnType = parInsn.getType();
-		switch (insnType) {
-			case CONST:
-				return parInsn.getArg(0);
-			case CONST_STR:
-				return ((ConstStringNode) parInsn).getString();
-			case CONST_CLASS:
-				return ((ConstClassNode) parInsn).getClsType();
-			case SGET:
-				FieldInfo f = (FieldInfo) ((IndexInsnNode) parInsn).getIndex();
-				FieldNode fieldNode = dex.resolveField(f);
-				if (fieldNode != null) {
-					FieldValueAttr attr = fieldNode.get(AType.FIELD_VALUE);
-					if (attr != null) {
-						return attr.getValue();
-					}
-				} else {
-					LOG.warn("Field {} not found in dex {}", f, dex);
-				}
-				break;
-		}
-		return null;
+		return InsnUtils.getConstValueByInsn(dex, parInsn);
 	}
 
 	@Override

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

@@ -318,36 +318,4 @@ public class CodeShrinker extends AbstractVisitor {
 		}
 		throw new JadxRuntimeException("Can't process instruction move : " + assignBlock);
 	}
-
-	@Deprecated
-	public static InsnArg inlineArgument(MethodNode mth, RegisterArg arg) {
-		InsnNode assignInsn = arg.getAssignInsn();
-		if (assignInsn == null) {
-			return null;
-		}
-		// recursively wrap all instructions
-		List<RegisterArg> list = new ArrayList<RegisterArg>();
-		List<RegisterArg> args = mth.getArguments(false);
-		int i = 0;
-		do {
-			list.clear();
-			assignInsn.getRegisterArgs(list);
-			for (RegisterArg rarg : list) {
-				InsnNode ai = rarg.getAssignInsn();
-				if (ai != assignInsn && ai != null && ai != rarg.getParentInsn()) {
-					inline(rarg, ai, null, mth);
-				}
-			}
-			// remove method args
-			if (!list.isEmpty() && !args.isEmpty()) {
-				list.removeAll(args);
-			}
-			i++;
-			if (i > 1000) {
-				throw new JadxRuntimeException("Can't inline arguments for: " + arg + " insn: " + assignInsn);
-			}
-		} while (!list.isEmpty());
-
-		return arg.wrapInstruction(assignInsn);
-	}
 }

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

@@ -11,24 +11,29 @@ import jadx.core.dex.instructions.IndexInsnNode;
 import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.instructions.args.ArgType;
 import jadx.core.dex.instructions.args.InsnArg;
-import jadx.core.dex.instructions.args.RegisterArg;
+import jadx.core.dex.instructions.args.InsnWrapArg;
 import jadx.core.dex.instructions.mods.ConstructorInsn;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.ClassNode;
+import jadx.core.dex.nodes.DexNode;
 import jadx.core.dex.nodes.FieldNode;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.utils.ErrorsCounter;
+import jadx.core.utils.InsnUtils;
 import jadx.core.utils.exceptions.JadxException;
 
 import java.util.ArrayList;
-import java.util.Iterator;
 import java.util.List;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-// TODO: run after code shrinker at final stage
+@JadxVisitor(
+		name = "EnumVisitor",
+		desc = "Restore enum classes",
+		runAfter = {CodeShrinker.class, ModVisitor.class}
+)
 public class EnumVisitor extends AbstractVisitor {
 	private static final Logger LOG = LoggerFactory.getLogger(EnumVisitor.class);
 
@@ -37,6 +42,25 @@ public class EnumVisitor extends AbstractVisitor {
 		if (!cls.isEnum()) {
 			return true;
 		}
+		// search class init method
+		MethodNode staticMethod = null;
+		for (MethodNode mth : cls.getMethods()) {
+			MethodInfo mi = mth.getMethodInfo();
+			if (mi.isClassInit()) {
+				staticMethod = mth;
+				break;
+			}
+		}
+		if (staticMethod == null) {
+			ErrorsCounter.classError(cls, "Enum class init method not found");
+			return true;
+		}
+
+		ArgType clsType = cls.getClassInfo().getType();
+		String enumConstructor = "<init>(Ljava/lang/String;I)V";
+		// TODO: detect these methods by analyzing method instructions
+		String valuesOfMethod = "valueOf(Ljava/lang/String;)" + TypeGen.signature(clsType);
+		String valuesMethod = "values()" + TypeGen.signature(ArgType.array(clsType));
 
 		// collect enum fields, remove synthetic
 		List<FieldNode> enumFields = new ArrayList<FieldNode>();
@@ -49,131 +73,133 @@ public class EnumVisitor extends AbstractVisitor {
 			}
 		}
 
-		MethodNode staticMethod = null;
-
-		ArgType clsType = cls.getClassInfo().getType();
-		String enumConstructor = "<init>(Ljava/lang/String;I)V";
-		String valuesOfMethod = "valueOf(Ljava/lang/String;)" + TypeGen.signature(clsType);
-		String valuesMethod = "values()" + TypeGen.signature(ArgType.array(clsType));
-
 		// remove synthetic methods
-		for (Iterator<MethodNode> it = cls.getMethods().iterator(); it.hasNext(); ) {
-			MethodNode mth = it.next();
+		for (MethodNode mth : cls.getMethods()) {
 			MethodInfo mi = mth.getMethodInfo();
 			if (mi.isClassInit()) {
-				staticMethod = mth;
-			} else {
-				String shortId = mi.getShortId();
-				boolean isSynthetic = mth.getAccessFlags().isSynthetic();
-				if (mi.isConstructor() && !isSynthetic) {
-					if (shortId.equals(enumConstructor)) {
-						it.remove();
-					}
-				} else if (isSynthetic
-						|| shortId.equals(valuesMethod)
-						|| shortId.equals(valuesOfMethod)) {
-					it.remove();
+				continue;
+			}
+			String shortId = mi.getShortId();
+			boolean isSynthetic = mth.getAccessFlags().isSynthetic();
+			if (mi.isConstructor() && !isSynthetic) {
+				if (shortId.equals(enumConstructor)) {
+					mth.add(AFlag.DONT_GENERATE);
 				}
+			} else if (isSynthetic
+					|| shortId.equals(valuesMethod)
+					|| shortId.equals(valuesOfMethod)) {
+				mth.add(AFlag.DONT_GENERATE);
 			}
 		}
 
 		EnumClassAttr attr = new EnumClassAttr(enumFields.size());
 		cls.addAttr(attr);
 
-		if (staticMethod == null) {
-			ErrorsCounter.classError(cls, "Enum class init method not found");
-			// for this broken enum puts found fields and mark as inconsistent
-			for (FieldNode field : enumFields) {
-				attr.getFields().add(new EnumField(field.getName(), 0));
-			}
-			return false;
-		}
 		attr.setStaticMethod(staticMethod);
+		ClassInfo classInfo = cls.getClassInfo();
 
 		// move enum specific instruction from static method to separate list
 		BlockNode staticBlock = staticMethod.getBasicBlocks().get(0);
-		ClassInfo classInfo = cls.getClassInfo();
-		List<InsnNode> insns = new ArrayList<InsnNode>();
+		List<InsnNode> enumPutInsns = new ArrayList<InsnNode>();
 		List<InsnNode> list = staticBlock.getInstructions();
 		int size = list.size();
 		for (int i = 0; i < size; i++) {
 			InsnNode insn = list.get(i);
-			insns.add(insn);
-			if (insn.getType() == InsnType.SPUT) {
-				IndexInsnNode fp = (IndexInsnNode) insn;
-				FieldInfo f = (FieldInfo) fp.getIndex();
-				if (f.getDeclClass().equals(classInfo)) {
-					FieldNode fieldNode = cls.searchField(f);
-					if (fieldNode != null
-							&& fieldNode.getAccessFlags().isSynthetic()
-							&& fieldNode.getType().isArray()
-							&& fieldNode.getType().getArrayRootElement().equals(classInfo.getType())) {
-						if (i == size - 1) {
-							cls.getMethods().remove(staticMethod);
-						} else {
-							list.subList(0, i + 1).clear();
-						}
-						break;
-					}
+			if (insn.getType() != InsnType.SPUT) {
+				continue;
+			}
+			FieldInfo f = (FieldInfo) ((IndexInsnNode) insn).getIndex();
+			if (!f.getDeclClass().equals(classInfo)) {
+				continue;
+			}
+			FieldNode fieldNode = cls.searchField(f);
+			if (fieldNode != null && isEnumArrayField(classInfo, fieldNode)) {
+				if (i == size - 1) {
+					staticMethod.add(AFlag.DONT_GENERATE);
+				} else {
+					list.subList(0, i + 1).clear();
 				}
+				break;
+			} else {
+				enumPutInsns.add(insn);
 			}
 		}
 
-		for (InsnNode insn : insns) {
-			if (insn.getType() == InsnType.CONSTRUCTOR) {
-				ConstructorInsn co = (ConstructorInsn) insn;
-				if (insn.getArgsCount() < 2) {
-					continue;
-				}
-				ClassInfo clsInfo = co.getClassType();
-				ClassNode constrCls = cls.dex().resolveClass(clsInfo);
-				if (constrCls == null) {
-					continue;
-				}
-				if (!clsInfo.equals(classInfo) && !constrCls.getAccessFlags().isEnum()) {
-					continue;
-				}
-				RegisterArg nameArg = (RegisterArg) insn.getArg(0);
-				// InsnArg pos = insn.getArg(1);
-				// TODO add check: pos == j
-				String name = (String) nameArg.getConstValue(cls.dex());
-				if (name == null) {
-					throw new JadxException("Unknown enum field name: " + cls);
-				}
-				EnumField field = new EnumField(name, insn.getArgsCount() - 2);
-				attr.getFields().add(field);
-				for (int i = 2; i < insn.getArgsCount(); i++) {
-					InsnArg iArg = insn.getArg(i);
-					InsnArg constrArg = iArg;
-					if (iArg.isLiteral()) {
-						constrArg = iArg;
-					} else if (iArg.isRegister()) {
-						constrArg = CodeShrinker.inlineArgument(staticMethod, (RegisterArg) iArg);
-						if (constrArg == null) {
-							throw new JadxException("Can't inline constructor arg in enum: " + cls);
-						}
-					}
-					field.getArgs().add(constrArg);
+		for (InsnNode putInsn : enumPutInsns) {
+			ConstructorInsn co = getConstructorInsn(putInsn);
+			if (co == null || co.getArgsCount() < 2) {
+				continue;
+			}
+			ClassInfo clsInfo = co.getClassType();
+			ClassNode constrCls = cls.dex().resolveClass(clsInfo);
+			if (constrCls == null) {
+				continue;
+			}
+			if (!clsInfo.equals(classInfo) && !constrCls.getAccessFlags().isEnum()) {
+				continue;
+			}
+			String name = getConstString(cls.dex(), co.getArg(0));
+			if (name == null) {
+				throw new JadxException("Unknown enum field name: " + cls);
+			}
+			EnumField field = new EnumField(name, co, 2);
+			attr.getFields().add(field);
+			if (!co.getClassType().equals(classInfo)) {
+				// enum contains additional methods
+				for (ClassNode innerCls : cls.getInnerClasses()) {
+					processEnumInnerCls(co, field, innerCls);
 				}
+			}
+		}
+		return false;
+	}
 
-				if (!co.getClassType().equals(classInfo)) {
-					// enum contains additional methods
-					for (ClassNode innerCls : cls.getInnerClasses()) {
-						if (innerCls.getClassInfo().equals(co.getClassType())) {
-							// remove constructor, because it is anonymous class
-							for (Iterator<?> mit = innerCls.getMethods().iterator(); mit.hasNext(); ) {
-								MethodNode innerMth = (MethodNode) mit.next();
-								if (innerMth.getAccessFlags().isConstructor()) {
-									mit.remove();
-								}
-							}
-							field.setCls(innerCls);
-							innerCls.add(AFlag.DONT_GENERATE);
-						}
-					}
-				}
+	private static void processEnumInnerCls(ConstructorInsn co, EnumField field, ClassNode innerCls) {
+		if (!innerCls.getClassInfo().equals(co.getClassType())) {
+			return;
+		}
+		// remove constructor, because it is anonymous class
+		for (MethodNode innerMth : innerCls.getMethods()) {
+			if (innerMth.getAccessFlags().isConstructor()) {
+				innerMth.add(AFlag.DONT_GENERATE);
+			}
+		}
+		field.setCls(innerCls);
+		innerCls.add(AFlag.DONT_GENERATE);
+	}
+
+	private boolean isEnumArrayField(ClassInfo classInfo, FieldNode fieldNode) {
+		if (fieldNode.getAccessFlags().isSynthetic()) {
+			ArgType fType = fieldNode.getType();
+			if (fType.isArray() && fType.getArrayRootElement().equals(classInfo.getType())) {
+				return true;
 			}
 		}
 		return false;
 	}
+
+	private ConstructorInsn getConstructorInsn(InsnNode putInsn) {
+		if (putInsn.getArgsCount() != 1) {
+			return null;
+		}
+		InsnArg arg = putInsn.getArg(0);
+		if (arg.isInsnWrap()) {
+			InsnNode wrapInsn = ((InsnWrapArg) arg).getWrapInsn();
+			if (wrapInsn.getType() == InsnType.CONSTRUCTOR) {
+				return (ConstructorInsn) wrapInsn;
+			}
+		}
+		return null;
+	}
+
+	private String getConstString(DexNode dex, InsnArg arg) {
+		if (arg.isInsnWrap()) {
+			InsnNode constInsn = ((InsnWrapArg) arg).getWrapInsn();
+			Object constValue = InsnUtils.getConstValueByInsn(dex, constInsn);
+			if (constValue instanceof String) {
+				return (String) constValue;
+			}
+		}
+		return null;
+	}
 }

jadx-core/src/main/java/jadx/core/utils/InsnUtils.java

 package jadx.core.utils;
 
+import jadx.core.dex.attributes.AType;
+import jadx.core.dex.info.FieldInfo;
+import jadx.core.dex.instructions.ConstClassNode;
+import jadx.core.dex.instructions.ConstStringNode;
+import jadx.core.dex.instructions.IndexInsnNode;
 import jadx.core.dex.instructions.InsnType;
+import jadx.core.dex.nodes.DexNode;
+import jadx.core.dex.nodes.FieldNode;
+import jadx.core.dex.nodes.InsnNode;
+import jadx.core.dex.nodes.parser.FieldValueAttr;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
+import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import com.android.dx.io.instructions.DecodedInstruction;
 
 public class InsnUtils {
 
+	private static final Logger LOG = LoggerFactory.getLogger(InsnUtils.class);
+
 	private InsnUtils() {
 	}
 
@@ -48,4 +63,34 @@ public class InsnUtils {
 			return index.toString();
 		}
 	}
+
+	/**
+	 * Return constant value from insn or null if not constant.
+	 *
+	 * @return LiteralArg, String, ArgType or null
+	 */
+	@Nullable
+	public static Object getConstValueByInsn(DexNode dex, InsnNode insn) {
+		switch (insn.getType()) {
+			case CONST:
+				return insn.getArg(0);
+			case CONST_STR:
+				return ((ConstStringNode) insn).getString();
+			case CONST_CLASS:
+				return ((ConstClassNode) insn).getClsType();
+			case SGET:
+				FieldInfo f = (FieldInfo) ((IndexInsnNode) insn).getIndex();
+				FieldNode fieldNode = dex.resolveField(f);
+				if (fieldNode != null) {
+					FieldValueAttr attr = fieldNode.get(AType.FIELD_VALUE);
+					if (attr != null) {
+						return attr.getValue();
+					}
+				} else {
+					LOG.warn("Field {} not found in dex {}", f, dex);
+				}
+				break;
+		}
+		return null;
+	}
 }

jadx-core/src/main/java/jadx/core/xmlgen/BinaryXMLParser.java

@@ -53,7 +53,7 @@ public class BinaryXMLParser extends CommonBinaryParser {
 	public BinaryXMLParser(RootNode root) {
 		try {
 			try {
-				Class rStyleCls = Class.forName(ANDROID_R_STYLE_CLS);
+				Class<?> rStyleCls = Class.forName(ANDROID_R_STYLE_CLS);
 				for (Field f : rStyleCls.getFields()) {
 					styleMap.put(f.getInt(f.getType()), f.getName());
 				}

jadx-core/src/test/java/jadx/tests/integration/enums/TestEnums4.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.containsOne;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+
+public class TestEnums4 extends IntegrationTest {
+
+	public static class TestCls {
+		public enum ResType {
+			CODE(".dex", ".class"),
+			MANIFEST("AndroidManifest.xml"),
+			XML(".xml"),
+			ARSC(".arsc"),
+			FONT(".ttf"),
+			IMG(".png", ".gif", ".jpg"),
+			LIB(".so"),
+			UNKNOWN;
+
+			private final String[] exts;
+
+			private ResType(String... exts) {
+				this.exts = exts;
+			}
+
+			public String[] getExts() {
+				return exts;
+			}
+		}
+
+		public void check() {
+			assertThat(ResType.CODE.getExts(), is(new String[]{".dex", ".class"}));
+		}
+	}
+
+	@Test
+	public void test() {
+		ClassNode cls = getClassNode(TestCls.class);
+		String code = cls.getCode().toString();
+
+		assertThat(code, containsOne("CODE(\".dex\", \".class\"),"));
+		assertThat(code, containsOne("ResType(String... exts) {"));
+//		assertThat(code, not(containsString("private ResType")));
+	}
+}