feat: add generic method information to .jcst (PR #564)

jadx-core/src/main/java/jadx/core/clsp/ClsSet.java

@@ -5,10 +5,12 @@ import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.File;
 import java.io.FileInputStream;
-import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -21,15 +23,15 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import jadx.core.dex.instructions.args.ArgType;
+import jadx.core.dex.instructions.args.RegisterArg;
 import jadx.core.dex.nodes.ClassNode;
+import jadx.core.dex.nodes.MethodNode;
 import jadx.core.dex.nodes.RootNode;
 import jadx.core.utils.exceptions.DecodeException;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 import jadx.core.utils.files.FileUtils;
 import jadx.core.utils.files.ZipSecurity;
 
-import static jadx.core.utils.files.FileUtils.close;
-
 /**
  * Classes list for import into classpath graph
  */
@@ -41,12 +43,14 @@ public class ClsSet {
 	private static final String CLST_PKG_PATH = ClsSet.class.getPackage().getName().replace('.', '/');
 
 	private static final String JADX_CLS_SET_HEADER = "jadx-cst";
-	private static final int VERSION = 1;
+	private static final int VERSION = 2;
 
 	private static final String STRING_CHARSET = "US-ASCII";
 
 	private static final NClass[] EMPTY_NCLASS_ARRAY = new NClass[0];
 
+	private enum ARG_TYPE {WILDCARD, GENERIC, GENERIC_TYPE, OBJECT, ARRAY, PRIMITIVE}
+
 	private NClass[] classes;
 
 	public void load(RootNode root) {
@@ -56,11 +60,13 @@ public class ClsSet {
 		for (ClassNode cls : list) {
 			String clsRawName = cls.getRawName();
 			if (cls.getAccessFlags().isPublic()) {
+				cls.load();
 				NClass nClass = new NClass(clsRawName, k);
 				if (names.put(clsRawName, nClass) != null) {
 					throw new JadxRuntimeException("Duplicate class: " + clsRawName);
 				}
 				k++;
+				nClass.setMethods(loadMethods(cls, nClass));
 			} else {
 				names.put(clsRawName, null);
 			}
@@ -80,6 +86,50 @@ public class ClsSet {
 		}
 	}
 
+	private NMethod[] loadMethods(ClassNode cls, NClass nClass) {
+		List<NMethod> methods = new ArrayList<>();
+		for (MethodNode m : cls.getMethods()) {
+			if (!m.getAccessFlags().isPublic()
+					&& !m.getAccessFlags().isProtected()) {
+				continue;
+			}
+
+			List<ArgType> args = new ArrayList<>();
+
+			boolean genericArg = false;
+			for (RegisterArg r: m.getArguments(false)) {
+				ArgType argType = r.getType();
+				if (argType.isGeneric()) {
+					args.add(argType);
+					genericArg = true;
+				} else if (argType.isGenericType()) {
+					args.add(argType);
+					genericArg = true;
+				} else {
+					args.add(null);
+				}
+			}
+
+			ArgType retType = m.getReturnType();
+			if (!retType.isGeneric() && !retType.isGenericType()) {
+				retType = null;
+			}
+
+			boolean varArgs = m.getAccessFlags().isVarArgs();
+
+			if (genericArg || retType != null || varArgs) {
+				methods.add(new NMethod(
+						m.getMethodInfo().getShortId(),
+						args.isEmpty()
+							? new ArgType[0]
+							: args.toArray(new ArgType[args.size()]),
+						retType,
+						varArgs));
+			}
+		}
+		return methods.toArray(new NMethod[methods.size()]);
+	}
+
 	public static NClass[] makeParentsArray(ClassNode cls, Map<String, NClass> names) {
 		List<NClass> parents = new ArrayList<>(1 + cls.getInterfaces().size());
 		ArgType superClass = cls.getSuperClass();
@@ -110,43 +160,129 @@ public class ClsSet {
 		return cls;
 	}
 
-	void save(File output) throws IOException {
-		FileUtils.makeDirsForFile(output);
-		try (BufferedOutputStream outputStream = new BufferedOutputStream(new FileOutputStream(output))) {
-			String outputName = output.getName();
-			if (outputName.endsWith(CLST_EXTENSION)) {
+	void save(Path path) throws IOException {
+		Files.createDirectories(path.getParent());
+		String outputName = path.getFileName().toString();
+		if (outputName.endsWith(CLST_EXTENSION)) {
+			try (BufferedOutputStream outputStream = new BufferedOutputStream(Files.newOutputStream(path))) {
 				save(outputStream);
-			} else if (outputName.endsWith(".jar")) {
-				ZipOutputStream out = new ZipOutputStream(outputStream);
-				try {
-					out.putNextEntry(new ZipEntry(CLST_PKG_PATH + '/' + CLST_FILENAME));
-					save(out);
-				} finally {
-					close(out);
+			}
+		} else if (outputName.endsWith(".jar")) {
+			Path temp = Files.createTempFile("jadx", ".zip");
+			Files.copy(path, temp, StandardCopyOption.REPLACE_EXISTING);
+
+			try (ZipOutputStream out = new ZipOutputStream(Files.newOutputStream(path));
+					ZipInputStream in = new ZipInputStream(Files.newInputStream(temp))) {
+				String clst = CLST_PKG_PATH + '/' + CLST_FILENAME;
+				out.putNextEntry(new ZipEntry(clst));
+				save(out);
+				ZipEntry entry = in.getNextEntry();
+				while (entry != null) {
+					if (!entry.getName().equals(clst)) {
+						out.putNextEntry(new ZipEntry(entry.getName()));
+						FileUtils.copyStream(in, out);
+					}
+					entry = in.getNextEntry();
 				}
-			} else {
-				throw new JadxRuntimeException("Unknown file format: " + outputName);
 			}
+			Files.delete(temp);
+
+		} else {
+			throw new JadxRuntimeException("Unknown file format: " + outputName);
 		}
 	}
 
 	public void save(OutputStream output) throws IOException {
-		try (DataOutputStream out = new DataOutputStream(output)) {
-			out.writeBytes(JADX_CLS_SET_HEADER);
-			out.writeByte(VERSION);
-
-			LOG.info("Classes count: {}", classes.length);
-			out.writeInt(classes.length);
-			for (NClass cls : classes) {
-				writeString(out, cls.getName());
-			}
-			for (NClass cls : classes) {
-				NClass[] parents = cls.getParents();
-				out.writeByte(parents.length);
-				for (NClass parent : parents) {
-					out.writeInt(parent.getId());
+		DataOutputStream out = new DataOutputStream(output);
+		out.writeBytes(JADX_CLS_SET_HEADER);
+		out.writeByte(VERSION);
+
+		LOG.info("Classes count: {}", classes.length);
+		Map<String, NClass> names = new HashMap<>(classes.length);
+		out.writeInt(classes.length);
+		for (NClass cls : classes) {
+			writeString(out, cls.getName());
+			names.put(cls.getName(), cls);
+		}
+		for (NClass cls : classes) {
+			NClass[] parents = cls.getParents();
+			out.writeByte(parents.length);
+			for (NClass parent : parents) {
+				out.writeInt(parent.getId());
+			}
+			NMethod[] methods = cls.getMethods();
+			out.writeByte(methods.length);
+			for (NMethod method : methods) {
+				writeMethod(out, method, names);
+			}
+		}
+	}
+
+	private static void writeMethod(DataOutputStream out, NMethod method, Map<String, NClass> names) throws IOException {
+		int argCount = 0;
+		ArgType[] argTypes = method.getArgType();
+		for (ArgType arg : argTypes) {
+			if (arg != null) {
+				argCount++;
+			}
+		}
+
+		writeLongString(out, method.getShortId());
+		out.writeByte(argCount);
+
+		// last argument first
+		for (int i = argTypes.length - 1; i >=0 ; i--) {
+			ArgType argType = argTypes[i];
+			if (argType != null) {
+				out.writeByte(i);
+				writeArgType(out, argType, names);
+			}
+		}
+
+		if (method.getReturnType() != null) {
+			out.writeBoolean(true);
+			writeArgType(out, method.getReturnType(), names);
+		} else {
+			out.writeBoolean(false);
+		}
+
+		out.writeBoolean(method.isVarArgs());
+	}
+
+	private static void writeArgType(DataOutputStream out, ArgType argType, Map<String, NClass> names) throws IOException {
+		if (argType.getWildcardType() != null) {
+			out.writeByte(ARG_TYPE.WILDCARD.ordinal());
+			int bounds = argType.getWildcardBounds();
+			out.writeByte(bounds);
+			if (bounds != 0) {
+				writeArgType(out, argType.getWildcardType(), names);
+			}
+		} else if (argType.isGeneric()) {
+			out.writeByte(ARG_TYPE.GENERIC.ordinal());
+			out.writeInt(names.get(argType.getObject()).getId());
+			ArgType[] types = argType.getGenericTypes();
+			if (types == null) {
+				out.writeByte(0);
+			} else {
+				out.writeByte(types.length);
+				for (ArgType type : types) {
+					writeArgType(out, type, names);
 				}
 			}
+		} else if (argType.isGenericType()) {
+			out.writeByte(ARG_TYPE.GENERIC_TYPE.ordinal());
+			writeString(out, argType.getObject());
+		} else if (argType.isObject()) {
+			out.writeByte(ARG_TYPE.OBJECT.ordinal());
+			out.writeInt(names.get(argType.getObject()).getId());
+		} else if (argType.isArray()) {
+			out.writeByte(ARG_TYPE.ARRAY.ordinal());
+			writeArgType(out, argType.getArrayElement(), names);
+		} else if (argType.isPrimitive()) {
+			out.writeByte(ARG_TYPE.PRIMITIVE.ordinal());
+			out.writeByte(argType.getPrimitiveType().getShortName().charAt(0));
+		} else {
+			throw new JadxRuntimeException("Cannot save type: " + argType);
 		}
 	}
 
@@ -203,18 +339,111 @@ public class ClsSet {
 					parents[j] = classes[in.readInt()];
 				}
 				classes[i].setParents(parents);
+
+				int mCount = in.readByte();
+				NMethod[] methods = new NMethod[mCount];
+				for (int j = 0; j < mCount; j++) {
+					methods[j] = readMethod(in);
+				}
+				classes[i].setMethods(methods);
 			}
 		}
 	}
 
-	private void writeString(DataOutputStream out, String name) throws IOException {
+	private NMethod readMethod(DataInputStream in) throws IOException {
+		String shortId = readLongString(in);
+		int argCount = in.readByte();
+		ArgType[] argTypes = null;
+		for (int i = 0; i < argCount; i++) {
+			int index = in.readByte();
+			ArgType argType = readArgType(in);
+			if (argTypes == null) {
+				argTypes = new ArgType[index + 1];
+			}
+			argTypes[index] = argType;
+		}
+		ArgType retType = in.readBoolean() ? readArgType(in) : null;
+		boolean varArgs = in.readBoolean();
+		return new NMethod(shortId, argTypes, retType, varArgs);
+	}
+
+	private ArgType readArgType(DataInputStream in) throws IOException {
+		int ordinal = in.readByte();
+		switch(ARG_TYPE.values()[ordinal]) {
+			case WILDCARD:
+				int bounds = in.readByte();
+				return bounds == 0
+						? ArgType.wildcard()
+						: ArgType.wildcard(readArgType(in), bounds);
+			case GENERIC:
+				String obj = classes[in.readInt()].getName();
+				int typeLength = in.readByte();
+				ArgType[] generics;
+				if (typeLength == 0) {
+					generics = null;
+				} else {
+					generics = new ArgType[typeLength];
+					for (int i = 0; i < typeLength; i++) {
+						generics[i] = readArgType(in);
+					}
+				}
+				return ArgType.generic(obj, generics);
+			case GENERIC_TYPE:
+				return ArgType.genericType(readString(in));
+			case OBJECT:
+				return ArgType.object(classes[in.readInt()].getName());
+			case ARRAY:
+				return ArgType.array(readArgType(in));
+			case PRIMITIVE:
+				int shortName = in.readByte();
+				switch(shortName) {
+				case 'Z':
+					return ArgType.BOOLEAN;
+				case 'C':
+					return ArgType.CHAR;
+				case 'B':
+					return ArgType.BYTE;
+				case 'S':
+					return ArgType.SHORT;
+				case 'I':
+					return ArgType.INT;
+				case 'F':
+					return ArgType.FLOAT;
+				case 'J':
+					return ArgType.LONG;
+				case 'D':
+					return ArgType.DOUBLE;
+				default:
+					return ArgType.VOID;
+				}
+			default:
+				throw new JadxRuntimeException("Unsupported Arg Type: " + ordinal);
+		}
+	}
+
+	private static void writeString(DataOutputStream out, String name) throws IOException {
 		byte[] bytes = name.getBytes(STRING_CHARSET);
 		out.writeByte(bytes.length);
 		out.write(bytes);
 	}
 
+	private static void writeLongString(DataOutputStream out, String name) throws IOException {
+		byte[] bytes = name.getBytes(STRING_CHARSET);
+		out.writeShort(bytes.length);
+		out.write(bytes);
+	}
+
 	private static String readString(DataInputStream in) throws IOException {
 		int len = in.readByte();
+		return readString(in, len);
+	}
+
+	private static String readLongString(DataInputStream in) throws IOException {
+		int len = in.readShort();
+		return readString(in, len);
+	}
+
+	private static String readString(DataInputStream in, int len) throws IOException {
 		byte[] bytes = new byte[len];
 		int count = in.read(bytes);
 		while (count != len) {

jadx-core/src/main/java/jadx/core/clsp/ConvertToClsSet.java

@@ -2,6 +2,8 @@ package jadx.core.clsp;
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -28,7 +30,7 @@ public class ConvertToClsSet {
 			usage();
 			System.exit(1);
 		}
-		File output = new File(args[0]);
+		Path output = Paths.get(args[0]);
 
 		List<InputFile> inputFiles = new ArrayList<>(args.length - 1);
 		for (int i = 1; i < args.length; i++) {

jadx-core/src/main/java/jadx/core/clsp/NClass.java

@@ -7,6 +7,7 @@ public class NClass {
 
 	private final String name;
 	private NClass[] parents;
+	private NMethod[] methods;
 	private final int id;
 
 	public NClass(String name, int id) {
@@ -51,4 +52,12 @@ public class NClass {
 	public String toString() {
 		return name;
 	}
+
+	public void setMethods(NMethod[] methods) {
+		this.methods = methods;
+	}
+
+	public NMethod[] getMethods() {
+		return methods;
+	}
 }

jadx-core/src/main/java/jadx/core/clsp/NMethod.java

+package jadx.core.clsp;
+
+import jadx.core.dex.instructions.args.ArgType;
+
+/**
+ * Generic method node in classpath graph.
+ */
+public class NMethod {
+
+	private final String shortId;
+	private final ArgType[] argType;
+	private final ArgType retType;
+	private final boolean varArgs;
+
+	public NMethod(String shortId, ArgType[] argType, ArgType retType, boolean varArgs) {
+		this.shortId = shortId;
+		this.argType = argType;
+		this.retType = retType;
+		this.varArgs = varArgs;
+	}
+
+	public String getShortId() {
+		return shortId;
+	}
+
+	public ArgType[] getArgType() {
+		return argType;
+	}
+
+	public ArgType getReturnType() {
+		return retType;
+	}
+
+	public boolean isVarArgs() {
+		return varArgs;
+	}
+}

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

@@ -75,8 +75,8 @@ public abstract class ArgType {
 		return new WildcardType(OBJECT, 0);
 	}
 
-	public static ArgType wildcard(ArgType obj, int bound) {
-		return new WildcardType(obj, bound);
+	public static ArgType wildcard(ArgType obj, int bounds) {
+		return new WildcardType(obj, bounds);
 	}
 
 	public static ArgType generic(String sign) {
@@ -214,10 +214,10 @@ public abstract class ArgType {
 		private final ArgType type;
 		private final int bounds;
 
-		public WildcardType(ArgType obj, int bound) {
+		public WildcardType(ArgType obj, int bounds) {
 			super(OBJECT.getObject());
 			this.type = obj;
-			this.bounds = bound;
+			this.bounds = bounds;
 		}
 
 		@Override

jadx-core/src/main/java/jadx/core/utils/files/InputFile.java

@@ -36,7 +36,7 @@ public class InputFile {
 
 	public static void addFilesFrom(File file, List<InputFile> list, boolean... skipSources) throws IOException, DecodeException {
 		InputFile inputFile = new InputFile(file);
-		inputFile.searchDexFiles(skipSources[0]);
+		inputFile.searchDexFiles(skipSources.length == 0 ? false : skipSources[0]);
 		list.add(inputFile);
 	}