refactor: add method info caching to speed up initial loading

2020-08-07 17:38:08 +01:00
parent 1560284831
commit 7fed5534eb
20 changed files with 233 additions and 89 deletions
@@ -24,6 +24,8 @@ import jadx.plugins.input.dex.sections.DexConsts;
 public class DexFileLoader {
 	private static final Logger LOG = LoggerFactory.getLogger(DexFileLoader.class);

+	private static int dexUniqId = 1;
+
 	public static List<DexReader> collectDexFiles(List<Path> pathsList) {
 		return pathsList.stream()
 				.map(Path::toFile)
@@ -52,7 +54,7 @@ public class DexFileLoader {
 			}
 			if (isStartWithBytes(magic, DexConsts.DEX_FILE_MAGIC)) {
 				in.reset();
-				DexReader dexReader = new DexReader(inputFileName, readAllBytes(in));
+				DexReader dexReader = new DexReader(getNextUniqId(), inputFileName, readAllBytes(in));
 				return Collections.singletonList(dexReader);
 			}
 			if (file != null && isStartWithBytes(magic, DexConsts.ZIP_FILE_MAGIC)) {
@@ -97,4 +99,16 @@ public class DexFileLoader {
 	private static byte[] readAllBytes(InputStream in) throws IOException {
 		return ByteStreams.toByteArray(in);
 	}
+
+	private static int getNextUniqId() {
+		dexUniqId++;
+		if (dexUniqId >= 0xFFFF) {
+			resetDexUniqId();
+		}
+		return dexUniqId;
+	}
+
+	public static void resetDexUniqId() {
+		dexUniqId = 1;
+	}
 }
@@ -11,6 +11,10 @@ import jadx.api.plugins.input.data.impl.EmptyLoadResult;

 public class DexInputPlugin implements JadxInputPlugin {

+	public DexInputPlugin() {
+		DexFileLoader.resetDexUniqId();
+	}
+
 	@Override
 	public JadxPluginInfo getPluginInfo() {
 		return new JadxPluginInfo("dex-input", "DexInput", "Load .dex and .apk files");
@@ -10,12 +10,13 @@ import jadx.plugins.input.dex.sections.SectionReader;
 import jadx.plugins.input.dex.sections.annotations.AnnotationsParser;

 public class DexReader {
-
+	private final int uniqId;
 	private final String inputFileName;
 	private final ByteBuffer buf;
 	private final DexHeader header;

-	public DexReader(String inputFileName, byte[] content) {
+	public DexReader(int uniqId, String inputFileName, byte[] content) {
+		this.uniqId = uniqId;
 		this.inputFileName = inputFileName;
 		this.buf = ByteBuffer.wrap(content);
 		this.header = new DexHeader(new SectionReader(this, 0));
@@ -48,6 +49,10 @@ public class DexReader {
 		return inputFileName;
 	}

+	public int getUniqId() {
+		return uniqId;
+	}
+
 	@Override
 	public String toString() {
 		return inputFileName;
@@ -3,7 +3,7 @@ package jadx.plugins.input.dex.insns;
 import org.jetbrains.annotations.Nullable;

 import jadx.api.plugins.input.data.IFieldData;
-import jadx.api.plugins.input.data.IMethodData;
+import jadx.api.plugins.input.data.IMethodRef;
 import jadx.api.plugins.input.insns.InsnData;
 import jadx.api.plugins.input.insns.InsnIndexType;
 import jadx.api.plugins.input.insns.Opcode;
@@ -106,8 +106,8 @@ public class DexInsnData implements InsnData {
 	}

 	@Override
-	public IMethodData getIndexAsMethod() {
-		return externalReader.getMethodData(index);
+	public IMethodRef getIndexAsMethod() {
+		return externalReader.getMethodRef(index);
 	}

 	@Nullable
@@ -127,10 +127,10 @@ public class DexClassData implements IClassData {

 	private void visitMethods(Consumer<IMethodData> mthConsumer, SectionReader data, int directMthCount, int virtualMthCount) {
 		DexMethodData methodData = new DexMethodData(annotationsParser);
+		methodData.setMethodRef(new DexMethodRef());
 		Map<Integer, Integer> annotationOffsetMap = annotationsParser.readMethodsAnnotationOffsetMap();
 		Map<Integer, Integer> paramsAnnOffsetMap = annotationsParser.readMethodParamsAnnRefOffsetMap();

-		methodData.setParentClassType(getType());
 		methodData.setDirect(true);
 		readMethods(mthConsumer, data, methodData, directMthCount, annotationOffsetMap, paramsAnnOffsetMap);
 		methodData.setDirect(false);
@@ -145,7 +145,10 @@ public class DexClassData implements IClassData {
 			mthIdx += data.readUleb128();
 			int accFlags = data.readUleb128();
 			int codeOff = data.readUleb128();
-			in.fillMethodData(methodData, mthIdx);
+
+			DexMethodRef methodRef = methodData.getMethodRef();
+			methodRef.reset();
+			in.initMethodRef(mthIdx, methodRef);
 			methodData.setAccessFlags(accFlags);
 			if (codeOff == 0) {
 				methodData.setCodeReader(null);
@@ -9,16 +9,13 @@ import jadx.api.plugins.input.data.IMethodData;
 import jadx.api.plugins.input.data.annotations.IAnnotation;
 import jadx.plugins.input.dex.sections.annotations.AnnotationsParser;
 import jadx.plugins.input.dex.smali.SmaliPrinter;
-import jadx.plugins.input.dex.utils.Utils;

 public class DexMethodData implements IMethodData {
 	@Nullable
 	private final AnnotationsParser annotationsParser;

-	private String parentClassType;
-	private String returnType;
-	private List<String> argTypes;
-	private String name;
+	private DexMethodRef methodRef;
+
 	private int accessFlags;
 	private boolean isDirect;
 	private int annotationsOffset;
@@ -32,21 +29,12 @@ public class DexMethodData implements IMethodData {
 	}

 	@Override
-	public String getParentClassType() {
-		return parentClassType;
+	public DexMethodRef getMethodRef() {
+		return methodRef;
 	}

-	public void setParentClassType(String parentClassType) {
-		this.parentClassType = parentClassType;
-	}
-
-	@Override
-	public String getName() {
-		return name;
-	}
-
-	public void setName(String name) {
-		this.name = name;
+	public void setMethodRef(DexMethodRef methodRef) {
+		this.methodRef = methodRef;
 	}

 	@Override
@@ -58,24 +46,6 @@ public class DexMethodData implements IMethodData {
 		this.accessFlags = accessFlags;
 	}

-	@Override
-	public String getReturnType() {
-		return returnType;
-	}
-
-	public void setReturnType(String returnType) {
-		this.returnType = returnType;
-	}
-
-	@Override
-	public List<String> getArgTypes() {
-		return argTypes;
-	}
-
-	public void setArgTypes(List<String> argTypes) {
-		this.argTypes = argTypes;
-	}
-
 	@Override
 	public boolean isDirect() {
 		return isDirect;
@@ -127,7 +97,6 @@ public class DexMethodData implements IMethodData {

 	@Override
 	public String toString() {
-		return getParentClassType() + "->" + getName()
-				+ '(' + Utils.listToStr(getArgTypes()) + ")" + getReturnType();
+		return getMethodRef().toString();
 	}
 }
@@ -0,0 +1,97 @@
+package jadx.plugins.input.dex.sections;
+
+import java.util.List;
+
+import jadx.api.plugins.input.data.IMethodRef;
+import jadx.plugins.input.dex.DexReader;
+import jadx.plugins.input.dex.utils.Utils;
+
+public class DexMethodRef implements IMethodRef {
+
+	private int uniqId;
+	private String name;
+	private String parentClassType;
+	private String returnType;
+	private List<String> argTypes;
+
+	// lazy loading info
+	private int dexIdx;
+	private SectionReader sectionReader;
+
+	public void initUniqId(DexReader dexReader, int idx) {
+		this.uniqId = (dexReader.getUniqId() & 0xFFFF) << 16 | (idx & 0xFFFF);
+	}
+
+	@Override
+	public void load() {
+		if (sectionReader != null) {
+			sectionReader.loadMethodRef(this, dexIdx);
+			sectionReader = null;
+		}
+	}
+
+	public void setDexIdx(int dexIdx) {
+		this.dexIdx = dexIdx;
+	}
+
+	public void setSectionReader(SectionReader sectionReader) {
+		this.sectionReader = sectionReader;
+	}
+
+	@Override
+	public int getUniqId() {
+		return uniqId;
+	}
+
+	public void reset() {
+		name = null;
+		parentClassType = null;
+		returnType = null;
+		argTypes = null;
+	}
+
+	@Override
+	public String getParentClassType() {
+		return parentClassType;
+	}
+
+	public void setParentClassType(String parentClassType) {
+		this.parentClassType = parentClassType;
+	}
+
+	@Override
+	public String getName() {
+		return name;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	@Override
+	public String getReturnType() {
+		return returnType;
+	}
+
+	public void setReturnType(String returnType) {
+		this.returnType = returnType;
+	}
+
+	@Override
+	public List<String> getArgTypes() {
+		return argTypes;
+	}
+
+	public void setArgTypes(List<String> argTypes) {
+		this.argTypes = argTypes;
+	}
+
+	@Override
+	public String toString() {
+		if (name == null) {
+			// not loaded
+			return Integer.toHexString(uniqId);
+		}
+		return getParentClassType() + "->" + getName() + '(' + Utils.listToStr(getArgTypes()) + ")" + getReturnType();
+	}
+}
@@ -10,7 +10,6 @@ import java.util.List;
 import org.jetbrains.annotations.Nullable;

 import jadx.api.plugins.input.data.IFieldData;
-import jadx.api.plugins.input.data.IMethodData;
 import jadx.plugins.input.dex.DexReader;
 import jadx.plugins.input.dex.utils.Leb128;
 import jadx.plugins.input.dex.utils.MUtf8;
@@ -175,23 +174,29 @@ public class SectionReader {
 		return classTypeIdx;
 	}

-	public IMethodData getMethodData(int idx) {
-		DexMethodData methodData = new DexMethodData(null);
-		int clsTypeIdx = fillMethodData(methodData, idx);
-		methodData.setParentClassType(getType(clsTypeIdx));
-		return methodData;
+	public DexMethodRef getMethodRef(int idx) {
+		DexMethodRef methodRef = new DexMethodRef();
+		initMethodRef(idx, methodRef);
+		return methodRef;
 	}

-	public int fillMethodData(DexMethodData methodData, int idx) {
-		int methodIdsOff = dexReader.getHeader().getMethodIdsOff();
+	public void initMethodRef(int idx, DexMethodRef methodRef) {
+		methodRef.initUniqId(dexReader, idx);
+		methodRef.setDexIdx(idx);
+		methodRef.setSectionReader(this);
+	}
+
+	public void loadMethodRef(DexMethodRef methodRef, int idx) {
+		DexHeader header = dexReader.getHeader();
+		int methodIdsOff = header.getMethodIdsOff();
 		absPos(methodIdsOff + idx * 8);
 		int classTypeIdx = readUShort();
 		int protoIdx = readUShort();
 		int nameIdx = readInt();

-		int protoIdsOff = dexReader.getHeader().getProtoIdsOff();
+		int protoIdsOff = header.getProtoIdsOff();
 		absPos(protoIdsOff + protoIdx * 12);
-		int shortyIdx = readInt();
+		skip(4); // shortyIdx
 		int returnTypeIdx = readInt();
 		int paramsOff = readInt();

@@ -201,18 +206,19 @@ public class SectionReader {
 		} else {
 			argTypes = absPos(paramsOff).readTypeList();
 		}
-		methodData.setName(getString(nameIdx));
-		methodData.setReturnType(getType(returnTypeIdx));
-		methodData.setArgTypes(argTypes);
-		return classTypeIdx;
+		methodRef.setParentClassType(getType(classTypeIdx));
+		methodRef.setName(getString(nameIdx));
+		methodRef.setReturnType(getType(returnTypeIdx));
+		methodRef.setArgTypes(argTypes);
 	}

 	public List<String> getMethodParamTypes(int idx) {
-		int methodIdsOff = dexReader.getHeader().getMethodIdsOff();
+		DexHeader header = dexReader.getHeader();
+		int methodIdsOff = header.getMethodIdsOff();
 		absPos(methodIdsOff + idx * 8 + 2);
 		int protoIdx = readUShort();

-		int protoIdsOff = dexReader.getHeader().getProtoIdsOff();
+		int protoIdsOff = header.getProtoIdsOff();
 		absPos(protoIdsOff + protoIdx * 12 + 8);
 		int paramsOff = readInt();

@@ -63,7 +63,7 @@ public class EncodedValueParser {
 				return new EncodedValue(EncodedType.ENCODED_TYPE, ext.getType(parseUnsignedInt(in, size)));

 			case ENCODED_METHOD:
-				return new EncodedValue(EncodedType.ENCODED_METHOD, ext.getMethodData(parseUnsignedInt(in, size)));
+				return new EncodedValue(EncodedType.ENCODED_METHOD, ext.getMethodRef(parseUnsignedInt(in, size)));

 			case ENCODED_FIELD:
 			case ENCODED_ENUM:
@@ -3,6 +3,7 @@ package jadx.plugins.input.dex.smali;
 import jadx.api.plugins.input.data.AccessFlags;
 import jadx.api.plugins.input.data.ICodeReader;
 import jadx.plugins.input.dex.sections.DexMethodData;
+import jadx.plugins.input.dex.sections.DexMethodRef;

 import static jadx.api.plugins.input.data.AccessFlagsScope.METHOD;

@@ -13,9 +14,12 @@ public class SmaliPrinter {
 		SmaliCodeWriter codeWriter = new SmaliCodeWriter();
 		codeWriter.startLine(".method ");
 		codeWriter.add(AccessFlags.format(mth.getAccessFlags(), METHOD));
-		codeWriter.add(mth.getName());
-		codeWriter.add('(').addArgs(mth.getArgTypes()).add(')');
-		codeWriter.add(mth.getReturnType());
+
+		DexMethodRef methodRef = mth.getMethodRef();
+		methodRef.load();
+		codeWriter.add(methodRef.getName());
+		codeWriter.add('(').addArgs(methodRef.getArgTypes()).add(')');
+		codeWriter.add(methodRef.getReturnType());
 		codeWriter.incIndent();

 		ICodeReader codeReader = mth.getCodeReader();
@@ -6,6 +6,9 @@ import java.util.List;
 public class Utils {

 	public static <T> String listToStr(List<T> collection) {
+		if (collection == null) {
+			return "null";
+		}
 		StringBuilder sb = new StringBuilder();
 		Iterator<T> it = collection.iterator();
 		if (it.hasNext()) {
@@ -7,16 +7,11 @@ import org.jetbrains.annotations.Nullable;
 import jadx.api.plugins.input.data.annotations.IAnnotation;

 public interface IMethodData {
-	String getParentClassType();

-	String getName();
+	IMethodRef getMethodRef();

 	int getAccessFlags();

-	String getReturnType();
-
-	List<String> getArgTypes();
-
 	boolean isDirect();

 	@Nullable
@@ -0,0 +1,21 @@
+package jadx.api.plugins.input.data;
+
+import java.util.List;
+
+public interface IMethodRef {
+
+	int getUniqId();
+
+	/**
+	 * Lazy loading for method info, until load() is called only getUniqId() can be used
+	 */
+	void load();
+
+	String getParentClassType();
+
+	String getName();
+
+	String getReturnType();
+
+	List<String> getArgTypes();
+}
@@ -3,7 +3,7 @@ package jadx.api.plugins.input.insns;
 import org.jetbrains.annotations.Nullable;

 import jadx.api.plugins.input.data.IFieldData;
-import jadx.api.plugins.input.data.IMethodData;
+import jadx.api.plugins.input.data.IMethodRef;
 import jadx.api.plugins.input.insns.custom.ICustomPayload;

 public interface InsnData {
@@ -34,7 +34,7 @@ public interface InsnData {

 	IFieldData getIndexAsField();

-	IMethodData getIndexAsMethod();
+	IMethodRef getIndexAsMethod();

 	@Nullable
 	ICustomPayload getPayload();