core: add cache for JavaNodes, fix definition annotations

jadx-core/src/main/java/jadx/api/CodePosition.java

@@ -2,24 +2,32 @@ package jadx.api;
 
 public final class CodePosition {
 
-	private final JavaClass cls;
+	private final JavaNode node;
 	private final int line;
 	private final int offset;
 
-	public CodePosition(JavaClass cls, int line, int offset) {
-		this.cls = cls;
+	public CodePosition(JavaNode node, int line, int offset) {
+		this.node = node;
 		this.line = line;
 		this.offset = offset;
 	}
 
 	public CodePosition(int line, int offset) {
-		this.cls = null;
+		this.node = null;
 		this.line = line;
 		this.offset = offset;
 	}
 
+	public JavaNode getNode() {
+		return node;
+	}
+
 	public JavaClass getJavaClass() {
-		return cls;
+		JavaClass parent = node.getDeclaringClass();
+		if (parent == null && node instanceof JavaClass) {
+			return (JavaClass) node;
+		}
+		return parent;
 	}
 
 	public int getLine() {
@@ -30,10 +38,6 @@ public final class CodePosition {
 		return offset;
 	}
 
-	public boolean isSet() {
-		return line != 0 || offset != 0;
-	}
-
 	@Override
 	public boolean equals(Object o) {
 		if (this == o) {
@@ -53,6 +57,6 @@ public final class CodePosition {
 
 	@Override
 	public String toString() {
-		return line + ":" + offset + (cls != null ? " " + cls : "");
+		return line + ":" + offset + (node != null ? " " + node : "");
 	}
 }

jadx-core/src/main/java/jadx/api/JadxDecompiler.java

@@ -5,6 +5,8 @@ import jadx.core.ProcessClass;
 import jadx.core.codegen.CodeGen;
 import jadx.core.codegen.CodeWriter;
 import jadx.core.dex.nodes.ClassNode;
+import jadx.core.dex.nodes.FieldNode;
+import jadx.core.dex.nodes.MethodNode;
 import jadx.core.dex.nodes.RootNode;
 import jadx.core.dex.visitors.IDexTreeVisitor;
 import jadx.core.dex.visitors.SaveCode;
@@ -62,6 +64,10 @@ public final class JadxDecompiler {
 
 	private BinaryXMLParser xmlParser;
 
+	private Map<ClassNode, JavaClass> classesMap = new HashMap<ClassNode, JavaClass>();
+	private Map<MethodNode, JavaMethod> methodsMap = new HashMap<MethodNode, JavaMethod>();
+	private Map<FieldNode, JavaField> fieldsMap = new HashMap<FieldNode, JavaField>();
+
 	public JadxDecompiler() {
 		this(new DefaultJadxArgs());
 	}
@@ -189,8 +195,11 @@ public final class JadxDecompiler {
 		if (classes == null) {
 			List<ClassNode> classNodeList = root.getClasses(false);
 			List<JavaClass> clsList = new ArrayList<JavaClass>(classNodeList.size());
+			classesMap.clear();
 			for (ClassNode classNode : classNodeList) {
-				clsList.add(new JavaClass(classNode, this));
+				JavaClass javaClass = new JavaClass(classNode, this);
+				clsList.add(javaClass);
+				classesMap.put(classNode, javaClass);
 			}
 			classes = Collections.unmodifiableList(clsList);
 		}
@@ -292,16 +301,16 @@ public final class JadxDecompiler {
 		return xmlParser;
 	}
 
-	JavaClass findJavaClass(ClassNode cls) {
-		if (cls == null) {
-			return null;
-		}
-		for (JavaClass javaClass : getClasses()) {
-			if (javaClass.getClassNode().equals(cls)) {
-				return javaClass;
-			}
-		}
-		return null;
+	Map<ClassNode, JavaClass> getClassesMap() {
+		return classesMap;
+	}
+
+	Map<MethodNode, JavaMethod> getMethodsMap() {
+		return methodsMap;
+	}
+
+	Map<FieldNode, JavaField> getFieldsMap() {
+		return fieldsMap;
 	}
 
 	public IJadxArgs getArgs() {

jadx-core/src/main/java/jadx/api/JavaClass.java

@@ -11,9 +11,12 @@ import jadx.core.dex.nodes.MethodNode;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import org.jetbrains.annotations.Nullable;
+
 public final class JavaClass implements JavaNode {
 
 	private final JadxDecompiler decompiler;
@@ -44,14 +47,14 @@ public final class JavaClass implements JavaNode {
 		if (code == null) {
 			decompile();
 			code = cls.getCode();
+			if (code == null) {
+				return "";
+			}
 		}
-		if (code == null) {
-			return "";
-		}
-		return code.toString();
+		return code.getCodeStr();
 	}
 
-	public void decompile() {
+	public synchronized void decompile() {
 		if (decompiler == null) {
 			return;
 		}
@@ -66,6 +69,7 @@ public final class JavaClass implements JavaNode {
 	}
 
 	private void load() {
+		JadxDecompiler rootDecompiler = getRootDecompiler();
 		int inClsCount = cls.getInnerClasses().size();
 		if (inClsCount != 0) {
 			List<JavaClass> list = new ArrayList<JavaClass>(inClsCount);
@@ -74,6 +78,7 @@ public final class JavaClass implements JavaNode {
 					JavaClass javaClass = new JavaClass(inner, this);
 					javaClass.load();
 					list.add(javaClass);
+					rootDecompiler.getClassesMap().put(inner, javaClass);
 				}
 			}
 			this.innerClasses = Collections.unmodifiableList(list);
@@ -84,7 +89,9 @@ public final class JavaClass implements JavaNode {
 			List<JavaField> flds = new ArrayList<JavaField>(fieldsCount);
 			for (FieldNode f : cls.getFields()) {
 				if (!f.contains(AFlag.DONT_GENERATE)) {
-					flds.add(new JavaField(f, this));
+					JavaField javaField = new JavaField(f, this);
+					flds.add(javaField);
+					rootDecompiler.getFieldsMap().put(f, javaField);
 				}
 			}
 			this.fields = Collections.unmodifiableList(flds);
@@ -95,7 +102,9 @@ public final class JavaClass implements JavaNode {
 			List<JavaMethod> mths = new ArrayList<JavaMethod>(methodsCount);
 			for (MethodNode m : cls.getMethods()) {
 				if (!m.contains(AFlag.DONT_GENERATE)) {
-					mths.add(new JavaMethod(this, m));
+					JavaMethod javaMethod = new JavaMethod(this, m);
+					mths.add(javaMethod);
+					rootDecompiler.getMethodsMap().put(m, javaMethod);
 				}
 			}
 			Collections.sort(mths, new Comparator<JavaMethod>() {
@@ -108,38 +117,81 @@ public final class JavaClass implements JavaNode {
 		}
 	}
 
+	private JadxDecompiler getRootDecompiler() {
+		if (parent != null) {
+			return parent.getRootDecompiler();
+		}
+		return decompiler;
+	}
+
 	private Map<CodePosition, Object> getCodeAnnotations() {
 		decompile();
 		return cls.getCode().getAnnotations();
 	}
 
-	public CodePosition getDefinitionPosition(int line, int offset) {
+	public Map<CodePosition, JavaNode> getUsageMap() {
 		Map<CodePosition, Object> map = getCodeAnnotations();
-		if (map.isEmpty()) {
-			return null;
+		if (map.isEmpty() || decompiler == null) {
+			return Collections.emptyMap();
 		}
-		Object obj = map.get(new CodePosition(line, offset));
+		Map<CodePosition, JavaNode> resultMap = new HashMap<CodePosition, JavaNode>(map.size());
+		for (Map.Entry<CodePosition, Object> entry : map.entrySet()) {
+			CodePosition codePosition = entry.getKey();
+			Object obj = entry.getValue();
+			if (obj instanceof LineAttrNode) {
+				JavaNode node = convertNode(obj);
+				if (node != null) {
+					resultMap.put(codePosition, node);
+				}
+			}
+		}
+		return resultMap;
+	}
+
+	@Nullable
+	private JavaNode convertNode(Object obj) {
 		if (!(obj instanceof LineAttrNode)) {
 			return null;
 		}
-		ClassNode clsNode = null;
 		if (obj instanceof ClassNode) {
-			clsNode = (ClassNode) obj;
-		} else if (obj instanceof MethodNode) {
-			clsNode = ((MethodNode) obj).getParentClass();
-		} else if (obj instanceof FieldNode) {
-			clsNode = ((FieldNode) obj).getParentClass();
+			return getRootDecompiler().getClassesMap().get(obj);
+		}
+		if (obj instanceof MethodNode) {
+			return getRootDecompiler().getMethodsMap().get(obj);
 		}
-		if (clsNode == null) {
+		if (obj instanceof FieldNode) {
+			return getRootDecompiler().getFieldsMap().get(obj);
+		}
+		return null;
+	}
+
+	@Nullable
+	public JavaNode getJavaNodeAtPosition(int line, int offset) {
+		Map<CodePosition, Object> map = getCodeAnnotations();
+		if (map.isEmpty()) {
+			return null;
+		}
+		Object obj = map.get(new CodePosition(line, offset));
+		if (obj == null) {
 			return null;
 		}
-		clsNode = clsNode.getTopParentClass();
-		JavaClass jCls = decompiler.findJavaClass(clsNode);
-		if (jCls == null) {
+		return convertNode(obj);
+	}
+
+	@Nullable
+	public CodePosition getDefinitionPosition(int line, int offset) {
+		JavaNode javaNode = getJavaNodeAtPosition(line, offset);
+		if (javaNode == null) {
 			return null;
 		}
+		return getDefinitionPosition(javaNode);
+	}
+
+	@Nullable
+	public CodePosition getDefinitionPosition(JavaNode javaNode) {
+		JavaClass jCls = javaNode.getTopParentClass();
 		jCls.decompile();
-		int defLine = ((LineAttrNode) obj).getDecompiledLine();
+		int defLine = javaNode.getDecompiledLine();
 		if (defLine == 0) {
 			return null;
 		}
@@ -170,6 +222,11 @@ public final class JavaClass implements JavaNode {
 		return parent;
 	}
 
+	@Override
+	public JavaClass getTopParentClass() {
+		return parent == null ? this : parent.getTopParentClass();
+	}
+
 	public AccessInfo getAccessInfo() {
 		return cls.getAccessFlags();
 	}

jadx-core/src/main/java/jadx/api/JavaField.java

@@ -29,6 +29,11 @@ public final class JavaField implements JavaNode {
 		return parent;
 	}
 
+	@Override
+	public JavaClass getTopParentClass() {
+		return parent.getTopParentClass();
+	}
+
 	public AccessInfo getAccessFlags() {
 		return field.getAccessFlags();
 	}
@@ -40,4 +45,19 @@ public final class JavaField implements JavaNode {
 	public int getDecompiledLine() {
 		return field.getDecompiledLine();
 	}
+
+	@Override
+	public int hashCode() {
+		return field.hashCode();
+	}
+
+	@Override
+	public boolean equals(Object o) {
+		return this == o || o instanceof JavaField && field.equals(((JavaField) o).field);
+	}
+
+	@Override
+	public String toString() {
+		return field.toString();
+	}
 }

jadx-core/src/main/java/jadx/api/JavaMethod.java

@@ -30,6 +30,11 @@ public final class JavaMethod implements JavaNode {
 		return parent;
 	}
 
+	@Override
+	public JavaClass getTopParentClass() {
+		return parent.getTopParentClass();
+	}
+
 	public AccessInfo getAccessFlags() {
 		return mth.getAccessFlags();
 	}
@@ -53,4 +58,19 @@ public final class JavaMethod implements JavaNode {
 	public int getDecompiledLine() {
 		return mth.getDecompiledLine();
 	}
+
+	@Override
+	public int hashCode() {
+		return mth.hashCode();
+	}
+
+	@Override
+	public boolean equals(Object o) {
+		return this == o || o instanceof JavaMethod && mth.equals(((JavaMethod) o).mth);
+	}
+
+	@Override
+	public String toString() {
+		return mth.toString();
+	}
 }

jadx-core/src/main/java/jadx/api/JavaNode.java

@@ -7,4 +7,8 @@ public interface JavaNode {
 	String getFullName();
 
 	JavaClass getDeclaringClass();
+
+	JavaClass getTopParentClass();
+
+	int getDecompiledLine();
 }

jadx-core/src/main/java/jadx/api/JavaPackage.java

@@ -34,6 +34,16 @@ public final class JavaPackage implements JavaNode, Comparable<JavaPackage> {
 	}
 
 	@Override
+	public JavaClass getTopParentClass() {
+		return null;
+	}
+
+	@Override
+	public int getDecompiledLine() {
+		return 0;
+	}
+
+	@Override
 	public int compareTo(@NotNull JavaPackage o) {
 		return name.compareTo(o.name);
 	}

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

@@ -148,6 +148,7 @@ public class ClassGen {
 		} else {
 			clsCode.add("class ");
 		}
+		clsCode.attachDefinition(cls);
 		clsCode.add(cls.getShortName());
 
 		addGenericMap(clsCode, cls.getGenericMap());
@@ -179,7 +180,6 @@ public class ClassGen {
 				clsCode.add(' ');
 			}
 		}
-		clsCode.attachDefinition(cls);
 	}
 
 	public boolean addGenericMap(CodeWriter code, Map<ArgType, List<ArgType>> gmap) {
@@ -340,6 +340,7 @@ public class ClassGen {
 			code.startLine(f.getAccessFlags().makeString());
 			useType(code, f.getType());
 			code.add(' ');
+			code.attachDefinition(f);
 			code.add(f.getAlias());
 			FieldInitAttr fv = f.get(AType.FIELD_INIT);
 			if (fv != null) {
@@ -356,7 +357,6 @@ public class ClassGen {
 				}
 			}
 			code.add(';');
-			code.attachDefinition(f);
 		}
 	}
 

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

@@ -12,6 +12,7 @@ import java.util.Iterator;
 import java.util.Map;
 import java.util.TreeMap;
 
+import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,7 +34,9 @@ public class CodeWriter {
 			INDENT + INDENT + INDENT + INDENT + INDENT,
 	};
 
-	private final StringBuilder buf = new StringBuilder();
+	private StringBuilder buf = new StringBuilder();
+	@Nullable
+	private String code;
 	private String indentStr;
 	private int indent;
 
@@ -113,7 +116,7 @@ public class CodeWriter {
 		}
 		line += code.line;
 		offset = code.offset;
-		buf.append(code);
+		buf.append(code.buf);
 		return this;
 	}
 
@@ -194,12 +197,13 @@ public class CodeWriter {
 		}
 	}
 
-	public Object attachDefinition(LineAttrNode obj) {
-		return attachAnnotation(new DefinitionWrapper(obj), new CodePosition(line, offset));
+	public void attachDefinition(LineAttrNode obj) {
+		attachAnnotation(obj);
+		attachAnnotation(new DefinitionWrapper(obj), new CodePosition(line, offset));
 	}
 
-	public Object attachAnnotation(Object obj) {
-		return attachAnnotation(obj, new CodePosition(line, offset + 1));
+	public void attachAnnotation(Object obj) {
+		attachAnnotation(obj, new CodePosition(line, offset + 1));
 	}
 
 	private Object attachAnnotation(Object obj, CodePosition pos) {
@@ -232,7 +236,11 @@ public class CodeWriter {
 	}
 
 	public void finish() {
+		removeFirstEmptyLine();
 		buf.trimToSize();
+		code = buf.toString();
+		buf = null;
+
 		Iterator<Map.Entry<CodePosition, Object>> it = annotations.entrySet().iterator();
 		while (it.hasNext()) {
 			Map.Entry<CodePosition, Object> entry = it.next();
@@ -245,28 +253,23 @@ public class CodeWriter {
 		}
 	}
 
-	private static String removeFirstEmptyLine(String str) {
-		if (str.startsWith(NL)) {
-			return str.substring(NL.length());
+	private void removeFirstEmptyLine() {
+		if (buf.indexOf(NL) == 0) {
+			buf.delete(0, NL.length());
 		}
-		return str;
 	}
 
-	public int length() {
+	public int bufLength() {
 		return buf.length();
 	}
 
-	public boolean isEmpty() {
-		return buf.length() == 0;
-	}
-
-	public boolean notEmpty() {
-		return buf.length() != 0;
+	public String getCodeStr() {
+		return code;
 	}
 
 	@Override
 	public String toString() {
-		return buf.toString();
+		return buf == null ? code : buf.toString();
 	}
 
 	public void save(File dir, String subDir, String fileName) {
@@ -278,6 +281,9 @@ public class CodeWriter {
 	}
 
 	public void save(File file) {
+		if (code == null) {
+			finish();
+		}
 		String name = file.getName();
 		if (name.length() > MAX_FILENAME_LENGTH) {
 			int dotIndex = name.indexOf('.');
@@ -294,8 +300,6 @@ public class CodeWriter {
 		try {
 			FileUtils.makeDirsForFile(file);
 			out = new PrintWriter(file, "UTF-8");
-			String code = buf.toString();
-			code = removeFirstEmptyLine(code);
 			out.println(code);
 		} catch (Exception e) {
 			LOG.error("Save file error", e);
@@ -305,21 +309,4 @@ public class CodeWriter {
 			}
 		}
 	}
-
-	@Override
-	public int hashCode() {
-		return buf.toString().hashCode();
-	}
-
-	@Override
-	public boolean equals(Object o) {
-		if (this == o) {
-			return true;
-		}
-		if (!(o instanceof CodeWriter)) {
-			return false;
-		}
-		CodeWriter that = (CodeWriter) o;
-		return buf.toString().equals(that.buf.toString());
-	}
 }

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

@@ -79,9 +79,9 @@ public class InsnGen {
 	}
 
 	public void addArgDot(CodeWriter code, InsnArg arg) throws CodegenException {
-		int len = code.length();
+		int len = code.bufLength();
 		addArg(code, arg, true);
-		if (len != code.length()) {
+		if (len != code.bufLength()) {
 			code.add('.');
 		}
 	}

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

@@ -57,8 +57,8 @@ public class MethodGen {
 
 	public boolean addDefinition(CodeWriter code) {
 		if (mth.getMethodInfo().isClassInit()) {
-			code.startLine("static");
 			code.attachDefinition(mth);
+			code.startLine("static");
 			return true;
 		}
 		if (mth.contains(AFlag.ANONYMOUS_CONSTRUCTOR)) {
@@ -87,10 +87,12 @@ public class MethodGen {
 			code.add(' ');
 		}
 		if (mth.getAccessFlags().isConstructor()) {
+			code.attachDefinition(mth);
 			code.add(classGen.getClassNode().getShortName()); // constructor
 		} else {
 			classGen.useType(code, mth.getReturnType());
 			code.add(' ');
+			code.attachDefinition(mth);
 			code.add(mth.getAlias());
 		}
 		code.add('(');
@@ -113,7 +115,6 @@ public class MethodGen {
 		code.add(')');
 
 		annotationGen.addThrows(mth, code);
-		code.attachDefinition(mth);
 		return true;
 	}
 

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

@@ -46,7 +46,7 @@ public class RegionUtils {
 			List<IContainer> blocks = region.getSubBlocks();
 			return !blocks.isEmpty() && hasExitEdge(blocks.get(blocks.size() - 1));
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -68,7 +68,7 @@ public class RegionUtils {
 			}
 			return getLastInsn(blocks.get(blocks.size() - 1));
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -85,7 +85,7 @@ public class RegionUtils {
 			return !blocks.isEmpty()
 					&& hasExitBlock(blocks.get(blocks.size() - 1));
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -112,7 +112,7 @@ public class RegionUtils {
 			}
 			return count;
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -132,7 +132,7 @@ public class RegionUtils {
 			}
 			return false;
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -145,7 +145,7 @@ public class RegionUtils {
 				getAllRegionBlocks(block, blocks);
 			}
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -161,7 +161,7 @@ public class RegionUtils {
 			}
 			return false;
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -245,7 +245,7 @@ public class RegionUtils {
 			}
 			return null;
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + container.getClass());
+			throw new JadxRuntimeException(unknownContainerType(container));
 		}
 	}
 
@@ -267,7 +267,7 @@ public class RegionUtils {
 			}
 			return true;
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + cont.getClass());
+			throw new JadxRuntimeException(unknownContainerType(cont));
 		}
 	}
 
@@ -288,8 +288,14 @@ public class RegionUtils {
 			}
 			return true;
 		} else {
-			throw new JadxRuntimeException("Unknown container type: " + cont.getClass());
+			throw new JadxRuntimeException(unknownContainerType(cont));
 		}
 	}
 
+	protected static String unknownContainerType(IContainer container) {
+		if (container == null) {
+			return "Null container variable";
+		}
+		return "Unknown container type: " + container.getClass();
+	}
 }